<commit_before>//! A stably addressed token buffer supporting efficient traversal based on a
//! cheaply copyable cursor.
//!
//! *This module is available only if Syn is built with the `"parsing"` feature.*

// This module is heavily commented as it contains most of the unsafe code in
// Syn, and caution should be used when editing it. The public-facing interface
// is 100% safe but the implementation is fragile internally.

#[cfg(all(
    not(all(target_arch = "wasm32", any(target_os = "unknown", target_os = "wasi"))),
    feature = "proc-macro"
))]
use crate::proc_macro as pm;
use crate::Lifetime;
use proc_macro2::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};
use std::marker::PhantomData;
use std::ptr;
use std::slice;

/// Internal type which is used instead of `TokenTree` to represent a token tree
/// within a `TokenBuffer`.
enum Entry {
    // Mimicking types from proc-macro.
    Group(Group, TokenBuffer),
    Ident(Ident),
    Punct(Punct),
    Literal(Literal),
    // End entries contain a raw pointer to the entry from the containing
    // token tree, or null if this is the outermost level.
    End(*const Entry),
}

/// A buffer that can be efficiently traversed multiple times, unlike
/// `TokenStream` which requires a deep copy in order to traverse more than
/// once.
///
/// *This type is available only if Syn is built with the `"parsing"` feature.*
pub struct TokenBuffer {
    // NOTE: Do not implement clone on this - there are raw pointers inside
    // these entries which will be messed up. Moving the `TokenBuffer` itself is
    // safe as the data pointed to won't be moved.
    ptr: *const Entry,
    len: usize,
}

impl Drop for TokenBuffer {
    fn drop(&mut self) {
        unsafe {
            let slice = slice::from_raw_parts_mut(self.ptr as *mut Entry, self.len);
            let _ = Box::from_raw(slice);
        }
    }
}

impl TokenBuffer {
    // NOTE: Do not mutate the Vec returned from this function once it returns;
    // the address of its backing memory must remain stable.
    fn inner_new(stream: TokenStream, up: *const Entry) -> TokenBuffer {
        // Build up the entries list, recording the locations of any Groups
        // in the list to be processed later.
        let mut entries = Vec::new();
        let mut groups = Vec::new();
        for tt in stream {
            match tt {
                TokenTree::Ident(sym) => {
                    entries.push(Entry::Ident(sym));
                }
                TokenTree::Punct(op) => {
                    entries.push(Entry::Punct(op));
                }
                TokenTree::Literal(l) => {
                    entries.push(Entry::Literal(l));
                }
                TokenTree::Group(g) => {
                    // Record the index of the interesting entry, and store an
                    // `End(null)` there temporarily.
                    groups.push((entries.len(), g));
                    entries.push(Entry::End(ptr::null()));
                }
            }
        }
        // Add an `End` entry to the end with a reference to the enclosing token
        // stream which was passed in.
        entries.push(Entry::End(up));

        // NOTE: This is done to ensure that we don't accidentally modify the
        // length of the backing buffer. The backing buffer must remain at a
        // constant address after this point, as we are going to store a raw
        // pointer into it.
        let entries = entries.into_boxed_slice();
        let len = entries.len();
        // Convert boxed slice into a pointer to the first element early, to
        // avoid invalidating pointers into this slice when we move the Box.
        // See https://github.com/rust-lang/unsafe-code-guidelines/issues/326
        let entries = Box::into_raw(entries) as *mut Entry;
        for (idx, group) in groups {
            // We know that this index refers to one of the temporary
            // `End(null)` entries, and we know that the last entry is
            // `End(up)`, so the next index is also valid.
            let group_up = unsafe { entries.add(idx + 1) };

            // The end entry stored at the end of this Entry::Group should
            // point to the Entry which follows the Group in the list.
            let inner = Self::inner_new(group.stream(), group_up);
            unsafe { *entries.add(idx) = Entry::Group(group, inner) };
        }

        TokenBuffer { ptr: entries, len }
    }

    /// Creates a `TokenBuffer` containing all the tokens from the input
    /// `proc_macro::TokenStream`.
    ///
    /// *This method is available only if Syn is built with both the `"parsing"` and
    /// `"proc-macro"` features.*
    #[cfg(all(
        not(all(target_arch = "wasm32", any(target_os = "unknown", target_os = "wasi"))),
        feature = "proc-macro"
    ))]
    pub fn new(stream: pm::TokenStream) -> Self {
        Self::new2(stream.into())
    }

    /// Creates a `TokenBuffer` containing all the tokens from the input
    /// `proc_macro2::TokenStream`.
    pub fn new2(stream: TokenStream) -> Self {
        Self::inner_new(stream, ptr::null())
    }

    /// Creates a cursor referencing the first token in the buffer and able to
    /// traverse until the end of the buffer.
    pub fn begin(&self) -> Cursor {
        unsafe { Cursor::create(self.ptr, self.ptr.add(self.len - 1)) }
    }
}

/// A cheaply copyable cursor into a `TokenBuffer`.
///
/// This cursor holds a shared reference into the immutable data which is used
/// internally to represent a `TokenStream`, and can be efficiently manipulated
/// and copied around.
///
/// An empty `Cursor` can be created directly, or one may create a `TokenBuffer`
/// object and get a cursor to its first token with `begin()`.
///
/// Two cursors are equal if they have the same location in the same input
/// stream, and have the same scope.
///
/// *This type is available only if Syn is built with the `"parsing"` feature.*
pub struct Cursor<'a> {
    // The current entry which the `Cursor` is pointing at.
    ptr: *const Entry,
    // This is the only `Entry::End(..)` object which this cursor is allowed to
    // point at. All other `End` objects are skipped over in `Cursor::create`.
    scope: *const Entry,
    // Cursor is covariant in 'a. This field ensures that our pointers are still
    // valid.
    marker: PhantomData<&'a Entry>,
}

impl<'a> Cursor<'a> {
    /// Creates a cursor referencing a static empty TokenStream.
    pub fn empty() -> Self {
        // It's safe in this situation for us to put an `Entry` object in global
        // storage, despite it not actually being safe to send across threads
        // (`Ident` is a reference into a thread-local table). This is because
        // this entry never includes a `Ident` object.
        //
        // This wrapper struct allows us to break the rules and put a `Sync`
        // object in global storage.
        struct UnsafeSyncEntry(Entry);
        unsafe impl Sync for UnsafeSyncEntry {}
        static EMPTY_ENTRY: UnsafeSyncEntry = UnsafeSyncEntry(Entry::End(0 as *const Entry));

        Cursor {
            ptr: &EMPTY_ENTRY.0,
            scope: &EMPTY_ENTRY.0,
            marker: PhantomData,
        }
    }

    /// This create method intelligently exits non-explicitly-entered
    /// `None`-delimited scopes when the cursor reaches the end of them,
    /// allowing for them to be treated transparently.
    unsafe fn create(mut ptr: *const Entry, scope: *const Entry) -> Self {
        // NOTE: If we're looking at a `End(..)`, we want to advance the cursor
        // past it, unless `ptr == scope`, which means that we're at the edge of
        // our cursor's scope. We should only have `ptr != scope` at the exit
        // from None-delimited groups entered with `ignore_none`.
        while let Entry::End(exit) = *ptr {
            if ptr == scope {
                break;
            }
            ptr = exit;
        }

        Cursor {
            ptr,
            scope,
            marker: PhantomData,
        }
    }

    /// Get the current entry.
    fn entry(self) -> &'a Entry {
        unsafe { &*self.ptr }
    }

    /// Bump the cursor to point at the next token after the current one. This
    /// is undefined behavior if the cursor is currently looking at an
    /// `Entry::End`.
    unsafe fn bump(self) -> Cursor<'a> {
        Cursor::create(self.ptr.offset(1), self.scope)
    }

    /// While the cursor is looking at a `None`-delimited group, move it to look
    /// at the first token inside instead. If the group is empty, this will move
    /// the cursor past the `None`-delimited group.
    ///
    /// WARNING: This mutates its argument.
    fn ignore_none(&mut self) {
        while let Entry::Group(group, buf) = self.entry() {
            if group.delimiter() == Delimiter::None {
                // NOTE: We call `Cursor::create` here to make sure that
                // situations where we should immediately exit the span after
                // entering it are handled correctly.
                unsafe {
                    *self = Cursor::create(buf.ptr, self.scope);
                }
            } else {
                break;
            }
        }
    }

    /// Checks whether the cursor is currently pointing at the end of its valid
    /// scope.
    pub fn eof(self) -> bool {
        // We're at eof if we're at the end of our scope.
        self.ptr == self.scope
    }

    /// If the cursor is pointing at a `Group` with the given delimiter, returns
    /// a cursor into that group and one pointing to the next `TokenTree`.
    pub fn group(mut self, delim: Delimiter) -> Option<(Cursor<'a>, Span, Cursor<'a>)> {
        // If we're not trying to enter a none-delimited group, we want to
        // ignore them. We have to make sure to _not_ ignore them when we want
        // to enter them, of course. For obvious reasons.
        if delim != Delimiter::None {
            self.ignore_none();
        }

        if let Entry::Group(group, buf) = self.entry() {
            if group.delimiter() == delim {
                return Some((buf.begin(), group.span(), unsafe { self.bump() }));
            }
        }

        None
    }

    /// If the cursor is pointing at a `Ident`, returns it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn ident(mut self) -> Option<(Ident, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Ident(ident) => Some((ident.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Punct`, returns it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn punct(mut self) -> Option<(Punct, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Punct(op) if op.as_char() != '\'' => Some((op.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Literal`, return it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn literal(mut self) -> Option<(Literal, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Literal(lit) => Some((lit.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Lifetime`, returns it along with a
    /// cursor pointing at the next `TokenTree`.
    pub fn lifetime(mut self) -> Option<(Lifetime, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Punct(op) if op.as_char() == '\'' && op.spacing() == Spacing::Joint => {
                let next = unsafe { self.bump() };
                match next.ident() {
                    Some((ident, rest)) => {
                        let lifetime = Lifetime {
                            apostrophe: op.span(),
                            ident,
                        };
                        Some((lifetime, rest))
                    }
                    None => None,
                }
            }
            _ => None,
        }
    }

    /// Copies all remaining tokens visible from this cursor into a
    /// `TokenStream`.
    pub fn token_stream(self) -> TokenStream {
        let mut tts = Vec::new();
        let mut cursor = self;
        while let Some((tt, rest)) = cursor.token_tree() {
            tts.push(tt);
            cursor = rest;
        }
        tts.into_iter().collect()
    }

    /// If the cursor is pointing at a `TokenTree`, returns it along with a
    /// cursor pointing at the next `TokenTree`.
    ///
    /// Returns `None` if the cursor has reached the end of its stream.
    ///
    /// This method does not treat `None`-delimited groups as transparent, and
    /// will return a `Group(None, ..)` if the cursor is looking at one.
    pub fn token_tree(self) -> Option<(TokenTree, Cursor<'a>)> {
        let tree = match self.entry() {
            Entry::Group(group, _) => group.clone().into(),
            Entry::Literal(lit) => lit.clone().into(),
            Entry::Ident(ident) => ident.clone().into(),
            Entry::Punct(op) => op.clone().into(),
            Entry::End(..) => return None,
        };

        Some((tree, unsafe { self.bump() }))
    }

    /// Returns the `Span` of the current token, or `Span::call_site()` if this
    /// cursor points to eof.
    pub fn span(self) -> Span {
        match self.entry() {
            Entry::Group(group, _) => group.span(),
            Entry::Literal(l) => l.span(),
            Entry::Ident(t) => t.span(),
            Entry::Punct(o) => o.span(),
            Entry::End(..) => Span::call_site(),
        }
    }

    /// Skip over the next token without cloning it. Returns `None` if this
    /// cursor points to eof.
    ///
    /// This method treats `'lifetimes` as a single token.
    pub(crate) fn skip(self) -> Option<Cursor<'a>> {
        match self.entry() {
            Entry::End(..) => None,

            // Treat lifetimes as a single tt for the purposes of 'skip'.
            Entry::Punct(op) if op.as_char() == '\'' && op.spacing() == Spacing::Joint => {
                let next = unsafe { self.bump() };
                match next.entry() {
                    Entry::Ident(_) => Some(unsafe { next.bump() }),
                    _ => Some(next),
                }
            }
            _ => Some(unsafe { self.bump() }),
        }
    }
}

impl<'a> Copy for Cursor<'a> {}

impl<'a> Clone for Cursor<'a> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<'a> Eq for Cursor<'a> {}

impl<'a> PartialEq for Cursor<'a> {
    fn eq(&self, other: &Self) -> bool {
        let Cursor { ptr, scope, marker } = self;
        let _ = marker;
        *ptr == other.ptr && *scope == other.scope
    }
}

pub(crate) fn same_scope(a: Cursor, b: Cursor) -> bool {
    a.scope == b.scope
}

pub(crate) fn open_span_of_group(cursor: Cursor) -> Span {
    match cursor.entry() {
        Entry::Group(group, _) => group.span_open(),
        _ => cursor.span(),
    }
}

pub(crate) fn close_span_of_group(cursor: Cursor) -> Span {
    match cursor.entry() {
        Entry::Group(group, _) => group.span_close(),
        _ => cursor.span(),
    }
}
<commit_msg>Skip drop of entry known to be Entry::End placeholder<commit_after>//! A stably addressed token buffer supporting efficient traversal based on a
//! cheaply copyable cursor.
//!
//! *This module is available only if Syn is built with the `"parsing"` feature.*

// This module is heavily commented as it contains most of the unsafe code in
// Syn, and caution should be used when editing it. The public-facing interface
// is 100% safe but the implementation is fragile internally.

#[cfg(all(
    not(all(target_arch = "wasm32", any(target_os = "unknown", target_os = "wasi"))),
    feature = "proc-macro"
))]
use crate::proc_macro as pm;
use crate::Lifetime;
use proc_macro2::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};
use std::marker::PhantomData;
use std::ptr;
use std::slice;

/// Internal type which is used instead of `TokenTree` to represent a token tree
/// within a `TokenBuffer`.
enum Entry {
    // Mimicking types from proc-macro.
    Group(Group, TokenBuffer),
    Ident(Ident),
    Punct(Punct),
    Literal(Literal),
    // End entries contain a raw pointer to the entry from the containing
    // token tree, or null if this is the outermost level.
    End(*const Entry),
}

/// A buffer that can be efficiently traversed multiple times, unlike
/// `TokenStream` which requires a deep copy in order to traverse more than
/// once.
///
/// *This type is available only if Syn is built with the `"parsing"` feature.*
pub struct TokenBuffer {
    // NOTE: Do not implement clone on this - there are raw pointers inside
    // these entries which will be messed up. Moving the `TokenBuffer` itself is
    // safe as the data pointed to won't be moved.
    ptr: *const Entry,
    len: usize,
}

impl Drop for TokenBuffer {
    fn drop(&mut self) {
        unsafe {
            let slice = slice::from_raw_parts_mut(self.ptr as *mut Entry, self.len);
            let _ = Box::from_raw(slice);
        }
    }
}

impl TokenBuffer {
    // NOTE: Do not mutate the Vec returned from this function once it returns;
    // the address of its backing memory must remain stable.
    fn inner_new(stream: TokenStream, up: *const Entry) -> TokenBuffer {
        // Build up the entries list, recording the locations of any Groups
        // in the list to be processed later.
        let mut entries = Vec::new();
        let mut groups = Vec::new();
        for tt in stream {
            match tt {
                TokenTree::Ident(sym) => {
                    entries.push(Entry::Ident(sym));
                }
                TokenTree::Punct(op) => {
                    entries.push(Entry::Punct(op));
                }
                TokenTree::Literal(l) => {
                    entries.push(Entry::Literal(l));
                }
                TokenTree::Group(g) => {
                    // Record the index of the interesting entry, and store an
                    // `End(null)` there temporarily.
                    groups.push((entries.len(), g));
                    entries.push(Entry::End(ptr::null()));
                }
            }
        }
        // Add an `End` entry to the end with a reference to the enclosing token
        // stream which was passed in.
        entries.push(Entry::End(up));

        // NOTE: This is done to ensure that we don't accidentally modify the
        // length of the backing buffer. The backing buffer must remain at a
        // constant address after this point, as we are going to store a raw
        // pointer into it.
        let entries = entries.into_boxed_slice();
        let len = entries.len();
        // Convert boxed slice into a pointer to the first element early, to
        // avoid invalidating pointers into this slice when we move the Box.
        // See https://github.com/rust-lang/unsafe-code-guidelines/issues/326
        let entries = Box::into_raw(entries) as *mut Entry;
        for (idx, group) in groups {
            // We know that this index refers to one of the temporary
            // `End(null)` entries, and we know that the last entry is
            // `End(up)`, so the next index is also valid.
            let group_up = unsafe { entries.add(idx + 1) };

            // The end entry stored at the end of this Entry::Group should
            // point to the Entry which follows the Group in the list.
            let inner = Self::inner_new(group.stream(), group_up);

            // No need to drop the Entry already at that position, because we
            // know it is the `Entry::End` variant and does not require drop.
            unsafe { entries.add(idx).write(Entry::Group(group, inner)) };
        }

        TokenBuffer { ptr: entries, len }
    }

    /// Creates a `TokenBuffer` containing all the tokens from the input
    /// `proc_macro::TokenStream`.
    ///
    /// *This method is available only if Syn is built with both the `"parsing"` and
    /// `"proc-macro"` features.*
    #[cfg(all(
        not(all(target_arch = "wasm32", any(target_os = "unknown", target_os = "wasi"))),
        feature = "proc-macro"
    ))]
    pub fn new(stream: pm::TokenStream) -> Self {
        Self::new2(stream.into())
    }

    /// Creates a `TokenBuffer` containing all the tokens from the input
    /// `proc_macro2::TokenStream`.
    pub fn new2(stream: TokenStream) -> Self {
        Self::inner_new(stream, ptr::null())
    }

    /// Creates a cursor referencing the first token in the buffer and able to
    /// traverse until the end of the buffer.
    pub fn begin(&self) -> Cursor {
        unsafe { Cursor::create(self.ptr, self.ptr.add(self.len - 1)) }
    }
}

/// A cheaply copyable cursor into a `TokenBuffer`.
///
/// This cursor holds a shared reference into the immutable data which is used
/// internally to represent a `TokenStream`, and can be efficiently manipulated
/// and copied around.
///
/// An empty `Cursor` can be created directly, or one may create a `TokenBuffer`
/// object and get a cursor to its first token with `begin()`.
///
/// Two cursors are equal if they have the same location in the same input
/// stream, and have the same scope.
///
/// *This type is available only if Syn is built with the `"parsing"` feature.*
pub struct Cursor<'a> {
    // The current entry which the `Cursor` is pointing at.
    ptr: *const Entry,
    // This is the only `Entry::End(..)` object which this cursor is allowed to
    // point at. All other `End` objects are skipped over in `Cursor::create`.
    scope: *const Entry,
    // Cursor is covariant in 'a. This field ensures that our pointers are still
    // valid.
    marker: PhantomData<&'a Entry>,
}

impl<'a> Cursor<'a> {
    /// Creates a cursor referencing a static empty TokenStream.
    pub fn empty() -> Self {
        // It's safe in this situation for us to put an `Entry` object in global
        // storage, despite it not actually being safe to send across threads
        // (`Ident` is a reference into a thread-local table). This is because
        // this entry never includes a `Ident` object.
        //
        // This wrapper struct allows us to break the rules and put a `Sync`
        // object in global storage.
        struct UnsafeSyncEntry(Entry);
        unsafe impl Sync for UnsafeSyncEntry {}
        static EMPTY_ENTRY: UnsafeSyncEntry = UnsafeSyncEntry(Entry::End(0 as *const Entry));

        Cursor {
            ptr: &EMPTY_ENTRY.0,
            scope: &EMPTY_ENTRY.0,
            marker: PhantomData,
        }
    }

    /// This create method intelligently exits non-explicitly-entered
    /// `None`-delimited scopes when the cursor reaches the end of them,
    /// allowing for them to be treated transparently.
    unsafe fn create(mut ptr: *const Entry, scope: *const Entry) -> Self {
        // NOTE: If we're looking at a `End(..)`, we want to advance the cursor
        // past it, unless `ptr == scope`, which means that we're at the edge of
        // our cursor's scope. We should only have `ptr != scope` at the exit
        // from None-delimited groups entered with `ignore_none`.
        while let Entry::End(exit) = *ptr {
            if ptr == scope {
                break;
            }
            ptr = exit;
        }

        Cursor {
            ptr,
            scope,
            marker: PhantomData,
        }
    }

    /// Get the current entry.
    fn entry(self) -> &'a Entry {
        unsafe { &*self.ptr }
    }

    /// Bump the cursor to point at the next token after the current one. This
    /// is undefined behavior if the cursor is currently looking at an
    /// `Entry::End`.
    unsafe fn bump(self) -> Cursor<'a> {
        Cursor::create(self.ptr.offset(1), self.scope)
    }

    /// While the cursor is looking at a `None`-delimited group, move it to look
    /// at the first token inside instead. If the group is empty, this will move
    /// the cursor past the `None`-delimited group.
    ///
    /// WARNING: This mutates its argument.
    fn ignore_none(&mut self) {
        while let Entry::Group(group, buf) = self.entry() {
            if group.delimiter() == Delimiter::None {
                // NOTE: We call `Cursor::create` here to make sure that
                // situations where we should immediately exit the span after
                // entering it are handled correctly.
                unsafe {
                    *self = Cursor::create(buf.ptr, self.scope);
                }
            } else {
                break;
            }
        }
    }

    /// Checks whether the cursor is currently pointing at the end of its valid
    /// scope.
    pub fn eof(self) -> bool {
        // We're at eof if we're at the end of our scope.
        self.ptr == self.scope
    }

    /// If the cursor is pointing at a `Group` with the given delimiter, returns
    /// a cursor into that group and one pointing to the next `TokenTree`.
    pub fn group(mut self, delim: Delimiter) -> Option<(Cursor<'a>, Span, Cursor<'a>)> {
        // If we're not trying to enter a none-delimited group, we want to
        // ignore them. We have to make sure to _not_ ignore them when we want
        // to enter them, of course. For obvious reasons.
        if delim != Delimiter::None {
            self.ignore_none();
        }

        if let Entry::Group(group, buf) = self.entry() {
            if group.delimiter() == delim {
                return Some((buf.begin(), group.span(), unsafe { self.bump() }));
            }
        }

        None
    }

    /// If the cursor is pointing at a `Ident`, returns it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn ident(mut self) -> Option<(Ident, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Ident(ident) => Some((ident.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Punct`, returns it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn punct(mut self) -> Option<(Punct, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Punct(op) if op.as_char() != '\'' => Some((op.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Literal`, return it along with a cursor
    /// pointing at the next `TokenTree`.
    pub fn literal(mut self) -> Option<(Literal, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Literal(lit) => Some((lit.clone(), unsafe { self.bump() })),
            _ => None,
        }
    }

    /// If the cursor is pointing at a `Lifetime`, returns it along with a
    /// cursor pointing at the next `TokenTree`.
    pub fn lifetime(mut self) -> Option<(Lifetime, Cursor<'a>)> {
        self.ignore_none();
        match self.entry() {
            Entry::Punct(op) if op.as_char() == '\'' && op.spacing() == Spacing::Joint => {
                let next = unsafe { self.bump() };
                match next.ident() {
                    Some((ident, rest)) => {
                        let lifetime = Lifetime {
                            apostrophe: op.span(),
                            ident,
                        };
                        Some((lifetime, rest))
                    }
                    None => None,
                }
            }
            _ => None,
        }
    }

    /// Copies all remaining tokens visible from this cursor into a
    /// `TokenStream`.
    pub fn token_stream(self) -> TokenStream {
        let mut tts = Vec::new();
        let mut cursor = self;
        while let Some((tt, rest)) = cursor.token_tree() {
            tts.push(tt);
            cursor = rest;
        }
        tts.into_iter().collect()
    }

    /// If the cursor is pointing at a `TokenTree`, returns it along with a
    /// cursor pointing at the next `TokenTree`.
    ///
    /// Returns `None` if the cursor has reached the end of its stream.
    ///
    /// This method does not treat `None`-delimited groups as transparent, and
    /// will return a `Group(None, ..)` if the cursor is looking at one.
    pub fn token_tree(self) -> Option<(TokenTree, Cursor<'a>)> {
        let tree = match self.entry() {
            Entry::Group(group, _) => group.clone().into(),
            Entry::Literal(lit) => lit.clone().into(),
            Entry::Ident(ident) => ident.clone().into(),
            Entry::Punct(op) => op.clone().into(),
            Entry::End(..) => return None,
        };

        Some((tree, unsafe { self.bump() }))
    }

    /// Returns the `Span` of the current token, or `Span::call_site()` if this
    /// cursor points to eof.
    pub fn span(self) -> Span {
        match self.entry() {
            Entry::Group(group, _) => group.span(),
            Entry::Literal(l) => l.span(),
            Entry::Ident(t) => t.span(),
            Entry::Punct(o) => o.span(),
            Entry::End(..) => Span::call_site(),
        }
    }

    /// Skip over the next token without cloning it. Returns `None` if this
    /// cursor points to eof.
    ///
    /// This method treats `'lifetimes` as a single token.
    pub(crate) fn skip(self) -> Option<Cursor<'a>> {
        match self.entry() {
            Entry::End(..) => None,

            // Treat lifetimes as a single tt for the purposes of 'skip'.
            Entry::Punct(op) if op.as_char() == '\'' && op.spacing() == Spacing::Joint => {
                let next = unsafe { self.bump() };
                match next.entry() {
                    Entry::Ident(_) => Some(unsafe { next.bump() }),
                    _ => Some(next),
                }
            }
            _ => Some(unsafe { self.bump() }),
        }
    }
}

impl<'a> Copy for Cursor<'a> {}

impl<'a> Clone for Cursor<'a> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<'a> Eq for Cursor<'a> {}

impl<'a> PartialEq for Cursor<'a> {
    fn eq(&self, other: &Self) -> bool {
        let Cursor { ptr, scope, marker } = self;
        let _ = marker;
        *ptr == other.ptr && *scope == other.scope
    }
}

pub(crate) fn same_scope(a: Cursor, b: Cursor) -> bool {
    a.scope == b.scope
}

pub(crate) fn open_span_of_group(cursor: Cursor) -> Span {
    match cursor.entry() {
        Entry::Group(group, _) => group.span_open(),
        _ => cursor.span(),
    }
}

pub(crate) fn close_span_of_group(cursor: Cursor) -> Span {
    match cursor.entry() {
        Entry::Group(group, _) => group.span_close(),
        _ => cursor.span(),
    }
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use capstone_rust::capstone;
use capstone_rust::capstone_sys::{x86_insn, x86_insn_group};
use std::fmt;
use super::address::Address;
use super::instruction::Instruction;

/// A representation of an eBPF instruction.
#[derive(Debug)]
pub struct CapstoneInstruction {
    insn: capstone::Instr,
}

fn is_group_match<F>(i: &capstone::Instr, predicate: F) -> bool
where
    F: Fn(&&u32) -> bool,
{
    match i.detail {
        None => false,
        Some(ref detail) => {
            match i.id {
                capstone::InstrIdArch::X86(_) => {
                    match detail.groups.iter().find(predicate) {
                        Some(_) => true,
                        None => false,
                    }
                }
                _ => unimplemented!(),
            }
        }
    }
}

impl Instruction for CapstoneInstruction {
    fn address(&self) -> Address {
        Address::new(self.insn.address)
    }

    fn comment(&self) -> Option<String> {
        None
    }

    fn mnemonic(&self) -> &str {
        &*self.insn.mnemonic
    }

    fn cycle_count(&self) -> Option<u32> {
        None
    }

    fn is_call(&self) -> bool {
        is_group_match(&self.insn, |&&x| x == x86_insn_group::X86_GRP_CALL.as_int())
    }

    fn is_local_conditional_jump(&self) -> bool {
        self.is_local_jump() &&
            match self.insn.id {
                capstone::InstrIdArch::X86(x86_insn::X86_INS_JMP) |
                capstone::InstrIdArch::X86(x86_insn::X86_INS_LOOP) |
                capstone::InstrIdArch::X86(x86_insn::X86_INS_LOOPE) |
                capstone::InstrIdArch::X86(x86_insn::X86_INS_LOOPNE) |
                capstone::InstrIdArch::X86(x86_insn::X86_INS_XBEGIN) |
                capstone::InstrIdArch::X86(_) => true,
                _ => unimplemented!(),
            }
    }

    fn is_local_jump(&self) -> bool {
        is_group_match(&self.insn, |&&x| x == x86_insn_group::X86_GRP_JUMP.as_int()) &&
            match self.insn.id {
                capstone::InstrIdArch::X86(x86_insn::X86_INS_LJMP) => false,
                capstone::InstrIdArch::X86(_) => true,
                _ => unimplemented!(),
            }
    }

    fn is_return(&self) -> bool {
        is_group_match(&self.insn, |&&x| x == x86_insn_group::X86_GRP_RET.as_int())
    }

    fn target_address(&self) -> Option<Address> {
        None
    }
}

impl fmt::Display for CapstoneInstruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} {}", self.insn.mnemonic, self.insn.op_str)
    }
}

#[cfg(test)]
mod tests {
    use capstone_rust::capstone as cs;
    use super::CapstoneInstruction;
    use super::super::{Address, Function, Symbol};

    #[test]
    fn test() {
        let code = vec![0x55, 0x48, 0x8b, 0x05, 0xb8, 0x13, 0x00, 0x00];

        let dec = cs::Capstone::new(cs::cs_arch::CS_ARCH_X86, cs::cs_mode::CS_MODE_32).unwrap();
        let buf = dec.disasm(code.as_slice(), 0, 0).unwrap();
        let is = buf.iter()
            .map(|insn| CapstoneInstruction { insn })
            .collect::<Vec<_>>();
        let f = Function::new(Symbol::new(Address::new(100000), Some("test")), is);

        assert!(f.control_flow_graph.entry_block.is_some());
    }
}
<commit_msg>capstone tests: Don't need a Vec for the code to disassemble.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use capstone_rust::capstone;
use capstone_rust::capstone_sys::{x86_insn, x86_insn_group};
use std::fmt;
use super::address::Address;
use super::instruction::Instruction;

/// A representation of an eBPF instruction.
#[derive(Debug)]
pub struct CapstoneInstruction {
    insn: capstone::Instr,
}

fn is_group_match<F>(i: &capstone::Instr, predicate: F) -> bool
where
    F: Fn(&&u32) -> bool,
{
    match i.detail {
        None => false,
        Some(ref detail) => {
            match i.id {
                capstone::InstrIdArch::X86(_) => {
                    match detail.groups.iter().find(predicate) {
                        Some(_) => true,
                        None => false,
                    }
                }
                _ => unimplemented!(),
            }
        }
    }
}

impl Instruction for CapstoneInstruction {
    fn address(&self) -> Address {
        Address::new(self.insn.address)
    }

    fn comment(&self) -> Option<String> {
        None
    }

    fn mnemonic(&self) -> &str {
        &*self.insn.mnemonic
    }

    fn cycle_count(&self) -> Option<u32> {
        None
    }

    fn is_call(&self) -> bool {
        is_group_match(&self.insn, |&&x| x == x86_insn_group::X86_GRP_CALL.as_int())
    }

    fn is_local_conditional_jump(&self) -> bool {
        self.is_local_jump() &&
            match self.insn.id {
                capstone::InstrIdArch::X86(x86_insn::X86_INS_JMP) |
                capstone::InstrIdArch::X86(x86_insn::X86_INS_LOOP) |
                capstone::InstrIdArch::X86(x86_insn::X86_INS_LOOPE) |
                capstone::InstrIdArch::X86(x86_insn::X86_INS_LOOPNE) |
                capstone::InstrIdArch::X86(x86_insn::X86_INS_XBEGIN) |
                capstone::InstrIdArch::X86(_) => true,
                _ => unimplemented!(),
            }
    }

    fn is_local_jump(&self) -> bool {
        is_group_match(&self.insn, |&&x| x == x86_insn_group::X86_GRP_JUMP.as_int()) &&
            match self.insn.id {
                capstone::InstrIdArch::X86(x86_insn::X86_INS_LJMP) => false,
                capstone::InstrIdArch::X86(_) => true,
                _ => unimplemented!(),
            }
    }

    fn is_return(&self) -> bool {
        is_group_match(&self.insn, |&&x| x == x86_insn_group::X86_GRP_RET.as_int())
    }

    fn target_address(&self) -> Option<Address> {
        None
    }
}

impl fmt::Display for CapstoneInstruction {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} {}", self.insn.mnemonic, self.insn.op_str)
    }
}

#[cfg(test)]
mod tests {
    use capstone_rust::capstone as cs;
    use super::CapstoneInstruction;
    use super::super::{Address, Function, Symbol};

    #[test]
    fn test() {
        let code = &[0x55, 0x48, 0x8b, 0x05, 0xb8, 0x13, 0x00, 0x00];

        let dec = cs::Capstone::new(cs::cs_arch::CS_ARCH_X86, cs::cs_mode::CS_MODE_32).unwrap();
        let buf = dec.disasm(code, 0, 0).unwrap();
        let is = buf.iter()
            .map(|insn| CapstoneInstruction { insn })
            .collect::<Vec<_>>();
        let f = Function::new(Symbol::new(Address::new(100000), Some("test")), is);

        assert!(f.control_flow_graph.entry_block.is_some());
    }
}
<|endoftext|>


<commit_before>use std::{fmt, cmp};

mod tests;

#[derive(Debug, PartialEq)]
pub struct Card {
    pub value: Value,
    pub suit: Suit
}

impl fmt::Display for Card {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}{}", self.value, self.suit)
    }
}

impl cmp::Ord for Card {
    fn cmp(&self, other: &Card) -> cmp::Ordering {
        self.value.cmp(&other.value)
    }
}

impl cmp::Eq for Card { }

impl cmp::PartialOrd for Card {
    fn partial_cmp(&self, other: &Card) -> Option<cmp::Ordering> {
        self.value.partial_cmp(&other.value)
    }
}

impl Card {
    pub fn new(value: Value, suit: Suit) -> Card {
        Card {
            value: value,
            suit: suit
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord)]
pub enum Value {
    None    =   0,
    AceLow  =   1,
    Two     =   2,
    Three   =   3,
    Four    =   4,
    Five    =   5,
    Six     =   6,
    Seven   =   7,
    Eight   =   8,
    Nine    =   9,
    Ten     =   10,
    Jack    =   11,
    Queen   =   12,
    King    =   13,
    Ace     =   14
}

impl fmt::Display for Value {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let value = match *self {
            Value::AceLow   =>  "1",
            Value::Two      =>  "2",
            Value::Three    =>  "3",
            Value::Four     =>  "4",
            Value::Five     =>  "5",
            Value::Six      =>  "6",
            Value::Seven    =>  "7",
            Value::Eight    =>  "8",
            Value::Nine     =>  "9",
            Value::Ten      =>  "T",
            Value::Jack     =>  "J",
            Value::Queen    =>  "Q",
            Value::King     =>  "K",
            Value::Ace      =>  "A",
            Value::None     =>  "-"
        };
        write!(f, "{}", value)
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
pub enum Suit {
    None,
    Diamonds,
    Clubs,
    Hearts,
    Spades
}

impl fmt::Display for Suit {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let suit = match *self {
            Suit::Diamonds  =>  "♢",
            Suit::Clubs     =>  "♣",
            Suit::Hearts    =>  "♡",
            Suit::Spades    =>  "♠",
            Suit::None      =>  "-"
        };
        write!(f, "{}", suit)
    }
}
<commit_msg>Removed AceLow card<commit_after>use std::{fmt, cmp};

mod tests;

#[derive(Debug, PartialEq)]
pub struct Card {
    pub value: Value,
    pub suit: Suit
}

impl fmt::Display for Card {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}{}", self.value, self.suit)
    }
}

impl cmp::Ord for Card {
    fn cmp(&self, other: &Card) -> cmp::Ordering {
        self.value.cmp(&other.value)
    }
}

impl cmp::Eq for Card { }

impl cmp::PartialOrd for Card {
    fn partial_cmp(&self, other: &Card) -> Option<cmp::Ordering> {
        self.value.partial_cmp(&other.value)
    }
}

impl Card {
    pub fn new(value: Value, suit: Suit) -> Card {
        Card {
            value: value,
            suit: suit
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord)]
pub enum Value {
    None    =   0,
    Two     =   2,
    Three   =   3,
    Four    =   4,
    Five    =   5,
    Six     =   6,
    Seven   =   7,
    Eight   =   8,
    Nine    =   9,
    Ten     =   10,
    Jack    =   11,
    Queen   =   12,
    King    =   13,
    Ace     =   14
}

impl fmt::Display for Value {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let value = match *self {
            Value::AceLow   =>  "1",
            Value::Two      =>  "2",
            Value::Three    =>  "3",
            Value::Four     =>  "4",
            Value::Five     =>  "5",
            Value::Six      =>  "6",
            Value::Seven    =>  "7",
            Value::Eight    =>  "8",
            Value::Nine     =>  "9",
            Value::Ten      =>  "T",
            Value::Jack     =>  "J",
            Value::Queen    =>  "Q",
            Value::King     =>  "K",
            Value::Ace      =>  "A",
            Value::None     =>  "-"
        };
        write!(f, "{}", value)
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
pub enum Suit {
    None,
    Diamonds,
    Clubs,
    Hearts,
    Spades
}

impl fmt::Display for Suit {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let suit = match *self {
            Suit::Diamonds  =>  "♢",
            Suit::Clubs     =>  "♣",
            Suit::Hearts    =>  "♡",
            Suit::Spades    =>  "♠",
            Suit::None      =>  "-"
        };
        write!(f, "{}", suit)
    }
}
<|endoftext|>


<commit_before>use std::io;
use byteorder::{ByteOrder, LittleEndian};
use key;

pub struct ChaCha20 {
    state: [u32; 16],
}

impl ChaCha20 {
    pub fn new(key: &[u8; 32]) -> io::Result<(Self, [u8; 12])> {
        let nonce = key::gen()?;
        let mut chacha20 = Self { state: [0; 16] };
        Self::setup_state(&mut chacha20.state, &key, &nonce);
        Ok((chacha20, nonce))
    }

    pub fn get_block(&self, counter: u32) -> [u8; 64] {
        Self::serialize_block(self.block(counter))
    }

    fn block(&self, counter: u32) -> [u32; 16] {
        let mut state = [0; 16];
        state.copy_from_slice(&self.state);
        state[12] = counter;
        for _ in 0..10 {
            Self::inner_block(&mut state);
        }
        for (x, &y) in state.iter_mut().zip(self.state.iter()) {
            *x = x.wrapping_add(y);
        }
        state[12] = state[12].wrapping_add(counter);
        state
    }

    fn setup_state(state: &mut [u32; 16], key: &[u8; 32], nonce: &[u8; 12]) {
        state[0] = 0x61707865;
        state[1] = 0x3320646e;
        state[2] = 0x79622d32;
        state[3] = 0x6b206574;
        state[4..12].copy_from_slice(&Self::convert_key(&key));
        state[13..].copy_from_slice(&Self::convert_nonce(&nonce));
    }

    fn inner_block(state: &mut [u32; 16]) {
        Self::quarter_round(state, 0, 4, 8, 12);
        Self::quarter_round(state, 1, 5, 9, 13);
        Self::quarter_round(state, 2, 6, 10, 14);
        Self::quarter_round(state, 3, 7, 11, 15);
        Self::quarter_round(state, 0, 5, 10, 15);
        Self::quarter_round(state, 1, 6, 11, 12);
        Self::quarter_round(state, 2, 7, 8, 13);
        Self::quarter_round(state, 3, 4, 9, 14);
    }

    fn quarter_round(state: &mut [u32; 16], i: usize, j: usize, k: usize, l: usize) {
        let mut a = state[i];
        let mut b = state[j];
        let mut c = state[k];
        let mut d = state[l];
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(16);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(12);
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(8);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(7);
        state[i] = a;
        state[j] = b;
        state[k] = c;
        state[l] = d;
    }

    fn convert_key(key: &[u8; 32]) -> [u32; 8] {
        let mut ret = [0; 8];
        for (byte, chunk) in ret.iter_mut().zip(key.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn convert_nonce(nonce: &[u8; 12]) -> [u32; 3] {
        let mut ret = [0; 3];
        for (byte, chunk) in ret.iter_mut().zip(nonce.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn serialize_block(block: [u32; 16]) -> [u8; 64] {
        let mut ret = [0; 64];
        for (chunk, &byte) in ret.chunks_mut(4).zip(block.iter()) {
            LittleEndian::write_u32(chunk, byte);
        }
        ret
    }
}

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

    const KEY: [u8; 32] = [
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
        0x05,
        0x06,
        0x07,
        0x08,
        0x09,
        0x0a,
        0x0b,
        0x0c,
        0x0d,
        0x0e,
        0x0f,
        0x10,
        0x11,
        0x12,
        0x13,
        0x14,
        0x15,
        0x16,
        0x17,
        0x18,
        0x19,
        0x1a,
        0x1b,
        0x1c,
        0x1d,
        0x1e,
        0x1f,
    ];
    const NONCE: [u8; 12] = [
        0x00,
        0x00,
        0x00,
        0x09,
        0x00,
        0x00,
        0x00,
        0x4a,
        0x00,
        0x00,
        0x00,
        0x00,
    ];

    #[test]
    fn test_new() {
        let (chacha20, nonce) = ChaCha20::new(&KEY).unwrap();
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &nonce);
        assert_eq!(state, chacha20.state);
    }

    #[test]
    fn test_get_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        let expected = [
            0x10,
            0xf1,
            0xe7,
            0xe4,
            0xd1,
            0x3b,
            0x59,
            0x15,
            0x50,
            0x0f,
            0xdd,
            0x1f,
            0xa3,
            0x20,
            0x71,
            0xc4,
            0xc7,
            0xd1,
            0xf4,
            0xc7,
            0x33,
            0xc0,
            0x68,
            0x03,
            0x04,
            0x22,
            0xaa,
            0x9a,
            0xc3,
            0xd4,
            0x6c,
            0x4e,
            0xd2,
            0x82,
            0x64,
            0x46,
            0x07,
            0x9f,
            0xaa,
            0x09,
            0x14,
            0xc2,
            0xd7,
            0x05,
            0xd9,
            0x8b,
            0x02,
            0xa2,
            0xb5,
            0x12,
            0x9c,
            0xd1,
            0xde,
            0x16,
            0x4e,
            0xb9,
            0xcb,
            0xd0,
            0x83,
            0xe8,
            0xa2,
            0x50,
            0x3c,
            0x4e,
        ];
        let block = chacha20.get_block(1);
        assert_eq!(expected.len(), block.len());
        for (lhs, rhs) in block.iter().zip(block.iter()) {
            assert_eq!(lhs, rhs);
        }
    }

    #[test]
    fn test_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        assert_eq!(
            [
                0xe4e7f110,
                0x15593bd1,
                0x1fdd0f50,
                0xc47120a3,
                0xc7f4d1c7,
                0x0368c033,
                0x9aaa2204,
                0x4e6cd4c3,
                0x466482d2,
                0x09aa9f07,
                0x05d7c214,
                0xa2028bd9,
                0xd19c12b5,
                0xb94e16de,
                0xe883d0cb,
                0x4e3c50a2,
            ],
            chacha20.block(1)
        );
    }

    #[test]
    fn test_setup_state() {
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &NONCE);
        assert_eq!(
            [
                0x61707865,
                0x3320646e,
                0x79622d32,
                0x6b206574,
                0x03020100,
                0x07060504,
                0x0b0a0908,
                0x0f0e0d0c,
                0x13121110,
                0x17161514,
                0x1b1a1918,
                0x1f1e1d1c,
                0x00000000,
                0x09000000,
                0x4a000000,
                0x00000000,
            ],
            state
        );
    }

    #[test]
    fn test_inner_block() {
        let mut state = [
            0x61707865,
            0x3320646e,
            0x79622d32,
            0x6b206574,
            0x03020100,
            0x07060504,
            0x0b0a0908,
            0x0f0e0d0c,
            0x13121110,
            0x17161514,
            0x1b1a1918,
            0x1f1e1d1c,
            0x00000001,
            0x09000000,
            0x4a000000,
            0x00000000,
        ];
        for _ in 0..10 {
            ChaCha20::inner_block(&mut state);
        }
        assert_eq!(
            [
                0x837778ab,
                0xe238d763,
                0xa67ae21e,
                0x5950bb2f,
                0xc4f2d0c7,
                0xfc62bb2f,
                0x8fa018fc,
                0x3f5ec7b7,
                0x335271c2,
                0xf29489f3,
                0xeabda8fc,
                0x82e46ebd,
                0xd19c12b4,
                0xb04e16de,
                0x9e83d0cb,
                0x4e3c50a2,
            ],
            state
        );
    }

    #[test]
    fn test_quarter_round() {
        let mut state = [0; 16];
        state[0] = 0x11111111;
        state[1] = 0x01020304;
        state[2] = 0x9b8d6f43;
        state[3] = 0x01234567;
        ChaCha20::quarter_round(&mut state, 0, 1, 2, 3);
        assert_eq!([0xea2a92f4, 0xcb1cf8ce, 0x4581472e, 0x5881c4bb], state[..4]);

        state = [
            0x879531e0,
            0xc5ecf37d,
            0x516461b1,
            0xc9a62f8a,
            0x44c20ef3,
            0x3390af7f,
            0xd9fc690b,
            0x2a5f714c,
            0x53372767,
            0xb00a5631,
            0x974c541a,
            0x359e9963,
            0x5c971061,
            0x3d631689,
            0x2098d9d6,
            0x91dbd320,
        ];
        ChaCha20::quarter_round(&mut state, 2, 7, 8, 13);
        assert_eq!(
            [
                0x879531e0,
                0xc5ecf37d,
                0xbdb886dc,
                0xc9a62f8a,
                0x44c20ef3,
                0x3390af7f,
                0xd9fc690b,
                0xcfacafd2,
                0xe46bea80,
                0xb00a5631,
                0x974c541a,
                0x359e9963,
                0x5c971061,
                0xccc07c79,
                0x2098d9d6,
                0x91dbd320,
            ],
            state
        );
    }

    #[test]
    fn test_convert_key() {
        assert_eq!(
            [
                0x03020100,
                0x07060504,
                0x0b0a0908,
                0x0f0e0d0c,
                0x13121110,
                0x17161514,
                0x1b1a1918,
                0x1f1e1d1c,
            ],
            ChaCha20::convert_key(&KEY)
        );
    }

    #[test]
    fn test_convert_nonce() {
        assert_eq!(
            [0x09000000, 0x4a000000, 0x00000000],
            ChaCha20::convert_nonce(&NONCE)
        );
    }

    #[test]
    fn test_serialize_block() {
        let expected = [
            0x10,
            0xf1,
            0xe7,
            0xe4,
            0xd1,
            0x3b,
            0x59,
            0x15,
            0x50,
            0x0f,
            0xdd,
            0x1f,
            0xa3,
            0x20,
            0x71,
            0xc4,
            0xc7,
            0xd1,
            0xf4,
            0xc7,
            0x33,
            0xc0,
            0x68,
            0x03,
            0x04,
            0x22,
            0xaa,
            0x9a,
            0xc3,
            0xd4,
            0x6c,
            0x4e,
            0xd2,
            0x82,
            0x64,
            0x46,
            0x07,
            0x9f,
            0xaa,
            0x09,
            0x14,
            0xc2,
            0xd7,
            0x05,
            0xd9,
            0x8b,
            0x02,
            0xa2,
            0xb5,
            0x12,
            0x9c,
            0xd1,
            0xde,
            0x16,
            0x4e,
            0xb9,
            0xcb,
            0xd0,
            0x83,
            0xe8,
            0xa2,
            0x50,
            0x3c,
            0x4e,
        ];
        let final_state = [
            0xe4e7f110,
            0x15593bd1,
            0x1fdd0f50,
            0xc47120a3,
            0xc7f4d1c7,
            0x0368c033,
            0x9aaa2204,
            0x4e6cd4c3,
            0x466482d2,
            0x09aa9f07,
            0x05d7c214,
            0xa2028bd9,
            0xd19c12b5,
            0xb94e16de,
            0xe883d0cb,
            0x4e3c50a2,
        ];
        let block = ChaCha20::serialize_block(final_state);
        assert_eq!(expected.len(), block.len());
        for (lhs, rhs) in expected.iter().zip(block.iter()) {
            assert_eq!(lhs, rhs);
        }
    }
}
<commit_msg>chacha20: deduplicate test code<commit_after>use std::io;
use byteorder::{ByteOrder, LittleEndian};
use key;

pub struct ChaCha20 {
    state: [u32; 16],
}

impl ChaCha20 {
    pub fn new(key: &[u8; 32]) -> io::Result<(Self, [u8; 12])> {
        let nonce = key::gen()?;
        let mut chacha20 = Self { state: [0; 16] };
        Self::setup_state(&mut chacha20.state, &key, &nonce);
        Ok((chacha20, nonce))
    }

    pub fn get_block(&self, counter: u32) -> [u8; 64] {
        Self::serialize_block(self.block(counter))
    }

    fn block(&self, counter: u32) -> [u32; 16] {
        let mut state = [0; 16];
        state.copy_from_slice(&self.state);
        state[12] = counter;
        for _ in 0..10 {
            Self::inner_block(&mut state);
        }
        for (x, &y) in state.iter_mut().zip(self.state.iter()) {
            *x = x.wrapping_add(y);
        }
        state[12] = state[12].wrapping_add(counter);
        state
    }

    fn setup_state(state: &mut [u32; 16], key: &[u8; 32], nonce: &[u8; 12]) {
        state[0] = 0x61707865;
        state[1] = 0x3320646e;
        state[2] = 0x79622d32;
        state[3] = 0x6b206574;
        state[4..12].copy_from_slice(&Self::convert_key(&key));
        state[13..].copy_from_slice(&Self::convert_nonce(&nonce));
    }

    fn inner_block(state: &mut [u32; 16]) {
        Self::quarter_round(state, 0, 4, 8, 12);
        Self::quarter_round(state, 1, 5, 9, 13);
        Self::quarter_round(state, 2, 6, 10, 14);
        Self::quarter_round(state, 3, 7, 11, 15);
        Self::quarter_round(state, 0, 5, 10, 15);
        Self::quarter_round(state, 1, 6, 11, 12);
        Self::quarter_round(state, 2, 7, 8, 13);
        Self::quarter_round(state, 3, 4, 9, 14);
    }

    fn quarter_round(state: &mut [u32; 16], i: usize, j: usize, k: usize, l: usize) {
        let mut a = state[i];
        let mut b = state[j];
        let mut c = state[k];
        let mut d = state[l];
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(16);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(12);
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(8);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(7);
        state[i] = a;
        state[j] = b;
        state[k] = c;
        state[l] = d;
    }

    fn convert_key(key: &[u8; 32]) -> [u32; 8] {
        let mut ret = [0; 8];
        for (byte, chunk) in ret.iter_mut().zip(key.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn convert_nonce(nonce: &[u8; 12]) -> [u32; 3] {
        let mut ret = [0; 3];
        for (byte, chunk) in ret.iter_mut().zip(nonce.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn serialize_block(block: [u32; 16]) -> [u8; 64] {
        let mut ret = [0; 64];
        for (chunk, &byte) in ret.chunks_mut(4).zip(block.iter()) {
            LittleEndian::write_u32(chunk, byte);
        }
        ret
    }
}

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

    const KEY: [u8; 32] = [
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
        0x05,
        0x06,
        0x07,
        0x08,
        0x09,
        0x0a,
        0x0b,
        0x0c,
        0x0d,
        0x0e,
        0x0f,
        0x10,
        0x11,
        0x12,
        0x13,
        0x14,
        0x15,
        0x16,
        0x17,
        0x18,
        0x19,
        0x1a,
        0x1b,
        0x1c,
        0x1d,
        0x1e,
        0x1f,
    ];
    const NONCE: [u8; 12] = [
        0x00,
        0x00,
        0x00,
        0x09,
        0x00,
        0x00,
        0x00,
        0x4a,
        0x00,
        0x00,
        0x00,
        0x00,
    ];
    const BLOCK_ONE: [u8; 64] = [
        0x10,
        0xf1,
        0xe7,
        0xe4,
        0xd1,
        0x3b,
        0x59,
        0x15,
        0x50,
        0x0f,
        0xdd,
        0x1f,
        0xa3,
        0x20,
        0x71,
        0xc4,
        0xc7,
        0xd1,
        0xf4,
        0xc7,
        0x33,
        0xc0,
        0x68,
        0x03,
        0x04,
        0x22,
        0xaa,
        0x9a,
        0xc3,
        0xd4,
        0x6c,
        0x4e,
        0xd2,
        0x82,
        0x64,
        0x46,
        0x07,
        0x9f,
        0xaa,
        0x09,
        0x14,
        0xc2,
        0xd7,
        0x05,
        0xd9,
        0x8b,
        0x02,
        0xa2,
        0xb5,
        0x12,
        0x9c,
        0xd1,
        0xde,
        0x16,
        0x4e,
        0xb9,
        0xcb,
        0xd0,
        0x83,
        0xe8,
        0xa2,
        0x50,
        0x3c,
        0x4e,
    ];
    const SETUP_STATE: [u32; 16] = [
        0x61707865,
        0x3320646e,
        0x79622d32,
        0x6b206574,
        0x03020100,
        0x07060504,
        0x0b0a0908,
        0x0f0e0d0c,
        0x13121110,
        0x17161514,
        0x1b1a1918,
        0x1f1e1d1c,
        0x00000000,
        0x09000000,
        0x4a000000,
        0x00000000,
    ];
    const FINAL_STATE: [u32; 16] = [
        0xe4e7f110,
        0x15593bd1,
        0x1fdd0f50,
        0xc47120a3,
        0xc7f4d1c7,
        0x0368c033,
        0x9aaa2204,
        0x4e6cd4c3,
        0x466482d2,
        0x09aa9f07,
        0x05d7c214,
        0xa2028bd9,
        0xd19c12b5,
        0xb94e16de,
        0xe883d0cb,
        0x4e3c50a2,
    ];

    #[test]
    fn test_new() {
        let (chacha20, nonce) = ChaCha20::new(&KEY).unwrap();
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &nonce);
        assert_eq!(state, chacha20.state);
    }

    fn check_serialized_block(block: &[u8]) {
        assert_eq!(BLOCK_ONE.len(), block.len());
        for (lhs, rhs) in BLOCK_ONE.iter().zip(block) {
            assert_eq!(lhs, rhs);
        }
    }

    #[test]
    fn test_get_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        check_serialized_block(&chacha20.get_block(1));
    }

    #[test]
    fn test_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        assert_eq!(FINAL_STATE, chacha20.block(1));
    }

    #[test]
    fn test_setup_state() {
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &NONCE);
        assert_eq!(SETUP_STATE, state);
    }

    #[test]
    fn test_inner_block() {
        let mut state = SETUP_STATE;
        state[12] = 1;
        for _ in 0..10 {
            ChaCha20::inner_block(&mut state);
        }
        assert_eq!(
            [
                0x837778ab,
                0xe238d763,
                0xa67ae21e,
                0x5950bb2f,
                0xc4f2d0c7,
                0xfc62bb2f,
                0x8fa018fc,
                0x3f5ec7b7,
                0x335271c2,
                0xf29489f3,
                0xeabda8fc,
                0x82e46ebd,
                0xd19c12b4,
                0xb04e16de,
                0x9e83d0cb,
                0x4e3c50a2,
            ],
            state
        );
    }

    #[test]
    fn test_quarter_round() {
        let mut state = [0; 16];
        state[0] = 0x11111111;
        state[1] = 0x01020304;
        state[2] = 0x9b8d6f43;
        state[3] = 0x01234567;
        ChaCha20::quarter_round(&mut state, 0, 1, 2, 3);
        assert_eq!([0xea2a92f4, 0xcb1cf8ce, 0x4581472e, 0x5881c4bb], state[..4]);

        state = [
            0x879531e0,
            0xc5ecf37d,
            0x516461b1,
            0xc9a62f8a,
            0x44c20ef3,
            0x3390af7f,
            0xd9fc690b,
            0x2a5f714c,
            0x53372767,
            0xb00a5631,
            0x974c541a,
            0x359e9963,
            0x5c971061,
            0x3d631689,
            0x2098d9d6,
            0x91dbd320,
        ];
        ChaCha20::quarter_round(&mut state, 2, 7, 8, 13);
        assert_eq!(
            [
                0x879531e0,
                0xc5ecf37d,
                0xbdb886dc,
                0xc9a62f8a,
                0x44c20ef3,
                0x3390af7f,
                0xd9fc690b,
                0xcfacafd2,
                0xe46bea80,
                0xb00a5631,
                0x974c541a,
                0x359e9963,
                0x5c971061,
                0xccc07c79,
                0x2098d9d6,
                0x91dbd320,
            ],
            state
        );
    }

    #[test]
    fn test_convert_key() {
        assert_eq!(
            [
                0x03020100,
                0x07060504,
                0x0b0a0908,
                0x0f0e0d0c,
                0x13121110,
                0x17161514,
                0x1b1a1918,
                0x1f1e1d1c,
            ],
            ChaCha20::convert_key(&KEY)
        );
    }

    #[test]
    fn test_convert_nonce() {
        assert_eq!(
            [0x09000000, 0x4a000000, 0x00000000],
            ChaCha20::convert_nonce(&NONCE)
        );
    }

    #[test]
    fn test_serialize_block() {
        check_serialized_block(&ChaCha20::serialize_block(FINAL_STATE));
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The ChaCha random number generator.

use std::num::Wrapping as w;
use {Rng, SeedableRng, Rand, w32};

const KEY_WORDS    : usize =  8; // 8 words for the 256-bit key
const STATE_WORDS  : usize = 16;
const CHACHA_ROUNDS: u32 = 20; // Cryptographically secure from 8 upwards as of this writing

/// A random number generator that uses the ChaCha20 algorithm [1].
///
/// The ChaCha algorithm is widely accepted as suitable for
/// cryptographic purposes, but this implementation has not been
/// verified as such. Prefer a generator like `OsRng` that defers to
/// the operating system for cases that need high security.
///
/// [1]: D. J. Bernstein, [*ChaCha, a variant of
/// Salsa20*](http://cr.yp.to/chacha.html)
#[derive(Copy, Clone)]
pub struct ChaChaRng {
    buffer:  [w32; STATE_WORDS], // Internal buffer of output
    state:   [w32; STATE_WORDS], // Initial state
    index:   usize,                 // Index into state
}

static EMPTY: ChaChaRng = ChaChaRng {
    buffer:  [w(0); STATE_WORDS],
    state:   [w(0); STATE_WORDS],
    index:   STATE_WORDS
};

macro_rules! quarter_round{
    ($a: expr, $b: expr, $c: expr, $d: expr) => {{
        $a = $a + $b; $d = $d ^ $a; $d = w($d.0.rotate_left(16));
        $c = $c + $d; $b = $b ^ $c; $b = w($b.0.rotate_left(12));
        $a = $a + $b; $d = $d ^ $a; $d = w($d.0.rotate_left( 8));
        $c = $c + $d; $b = $b ^ $c; $b = w($b.0.rotate_left( 7));
    }}
}

macro_rules! double_round{
    ($x: expr) => {{
        // Column round
        quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
        quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
        quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
        quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
        // Diagonal round
        quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
        quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
        quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
        quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
    }}
}

#[inline]
fn core(output: &mut [w32; STATE_WORDS], input: &[w32; STATE_WORDS]) {
    *output = *input;

    for _ in 0..CHACHA_ROUNDS / 2 {
        double_round!(output);
    }

    for i in 0..STATE_WORDS {
        output[i] = output[i] + input[i];
    }
}

impl ChaChaRng {

    /// Create an ChaCha random number generator using the default
    /// fixed key of 8 zero words.
    pub fn new_unseeded() -> ChaChaRng {
        let mut rng = EMPTY;
        rng.init(&[0; KEY_WORDS]);
        rng
    }

    /// Sets the internal 128-bit ChaCha counter to
    /// a user-provided value. This permits jumping
    /// arbitrarily ahead (or backwards) in the pseudorandom stream.
    ///
    /// Since the nonce words are used to extend the counter to 128 bits,
    /// users wishing to obtain the conventional ChaCha pseudorandom stream
    /// associated with a particular nonce can call this function with
    /// arguments `0, desired_nonce`.
    pub fn set_counter(&mut self, counter_low: u64, counter_high: u64) {
        self.state[12] = w((counter_low >>  0) as u32);
        self.state[13] = w((counter_low >> 32) as u32);
        self.state[14] = w((counter_high >>  0) as u32);
        self.state[15] = w((counter_high >> 32) as u32);
        self.index = STATE_WORDS; // force recomputation
    }

    /// Initializes `self.state` with the appropriate key and constants
    ///
    /// We deviate slightly from the ChaCha specification regarding
    /// the nonce, which is used to extend the counter to 128 bits.
    /// This is provably as strong as the original cipher, though,
    /// since any distinguishing attack on our variant also works
    /// against ChaCha with a chosen-nonce. See the XSalsa20 [1]
    /// security proof for a more involved example of this.
    ///
    /// The modified word layout is:
    /// ```text
    /// constant constant constant constant
    /// key      key      key      key
    /// key      key      key      key
    /// counter  counter  counter  counter
    /// ```
    /// [1]: Daniel J. Bernstein. [*Extending the Salsa20
    /// nonce.*](http://cr.yp.to/papers.html#xsalsa)
    fn init(&mut self, key: &[u32; KEY_WORDS]) {
        self.state[0] = w(0x61707865);
        self.state[1] = w(0x3320646E);
        self.state[2] = w(0x79622D32);
        self.state[3] = w(0x6B206574);

        for i in 0..KEY_WORDS {
            self.state[4+i] = w(key[i]);
        }

        self.state[12] = w(0);
        self.state[13] = w(0);
        self.state[14] = w(0);
        self.state[15] = w(0);

        self.index = STATE_WORDS;
    }

    /// Refill the internal output buffer (`self.buffer`)
    fn update(&mut self) {
        core(&mut self.buffer, &self.state);
        self.index = 0;
        // update 128-bit counter
        self.state[12] = self.state[12] + w(1);
        if self.state[12] != w(0) { return };
        self.state[13] = self.state[13] + w(1);
        if self.state[13] != w(0) { return };
        self.state[14] = self.state[14] + w(1);
        if self.state[14] != w(0) { return };
        self.state[15] = self.state[15] + w(1);
    }
}

impl Rng for ChaChaRng {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        if self.index == STATE_WORDS {
            self.update();
        }

        let value = self.buffer[self.index % STATE_WORDS];
        self.index += 1;
        value.0
    }
}

impl<'a> SeedableRng<&'a [u32]> for ChaChaRng {

    fn reseed(&mut self, seed: &'a [u32]) {
        // reset state
        self.init(&[0u32; KEY_WORDS]);
        // set key in place
        let key = &mut self.state[4 .. 4+KEY_WORDS];
        for (k, s) in key.iter_mut().zip(seed.iter()) {
            *k = w(*s);
        }
    }

    /// Create a ChaCha generator from a seed,
    /// obtained from a variable-length u32 array.
    /// Only up to 8 words are used; if less than 8
    /// words are used, the remaining are set to zero.
    fn from_seed(seed: &'a [u32]) -> ChaChaRng {
        let mut rng = EMPTY;
        rng.reseed(seed);
        rng
    }
}

impl Rand for ChaChaRng {
    fn rand<R: Rng>(other: &mut R) -> ChaChaRng {
        let mut key : [u32; KEY_WORDS] = [0; KEY_WORDS];
        for word in key.iter_mut() {
            *word = other.gen();
        }
        SeedableRng::from_seed(&key[..])
    }
}

#[cfg(test)]
mod test {
    use {Rng, SeedableRng};
    use super::ChaChaRng;

    #[test]
    fn test_rng_rand_seeded() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut ra: ChaChaRng = SeedableRng::from_seed(&s[..]);
        let mut rb: ChaChaRng = SeedableRng::from_seed(&s[..]);
        assert!(::test::iter_eq(ra.gen_ascii_chars().take(100),
                                rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_seeded() {
        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);
        let mut rb: ChaChaRng = SeedableRng::from_seed(seed);
        assert!(::test::iter_eq(ra.gen_ascii_chars().take(100),
                                rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_reseed() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut r: ChaChaRng = SeedableRng::from_seed(&s[..]);
        let string1: String = r.gen_ascii_chars().take(100).collect();

        r.reseed(&s);

        let string2: String = r.gen_ascii_chars().take(100).collect();
        assert_eq!(string1, string2);
    }

    #[test]
    fn test_rng_true_values() {
        // Test vectors 1 and 2 from
        // http://tools.ietf.org/html/draft-nir-cfrg-chacha20-poly1305-04
        let seed : &[_] = &[0u32; 8];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        let v = (0..16).map(|_| ra.next_u32()).collect::<Vec<_>>();
        assert_eq!(v,
                   vec!(0xade0b876, 0x903df1a0, 0xe56a5d40, 0x28bd8653,
                        0xb819d2bd, 0x1aed8da0, 0xccef36a8, 0xc70d778b,
                        0x7c5941da, 0x8d485751, 0x3fe02477, 0x374ad8b8,
                        0xf4b8436a, 0x1ca11815, 0x69b687c3, 0x8665eeb2));

        let v = (0..16).map(|_| ra.next_u32()).collect::<Vec<_>>();
        assert_eq!(v,
                   vec!(0xbee7079f, 0x7a385155, 0x7c97ba98, 0x0d082d73,
                        0xa0290fcb, 0x6965e348, 0x3e53c612, 0xed7aee32,
                        0x7621b729, 0x434ee69c, 0xb03371d5, 0xd539d874,
                        0x281fed31, 0x45fb0a51, 0x1f0ae1ac, 0x6f4d794b));

        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        // Store the 17*i-th 32-bit word,
        // i.e., the i-th word of the i-th 16-word block
        let mut v : Vec<u32> = Vec::new();
        for _ in 0..16 {
            v.push(ra.next_u32());
            for _ in 0..16 {
                ra.next_u32();
            }
        }

        assert_eq!(v,
                   vec!(0xf225c81a, 0x6ab1be57, 0x04d42951, 0x70858036,
                        0x49884684, 0x64efec72, 0x4be2d186, 0x3615b384,
                        0x11cfa18e, 0xd3c50049, 0x75c775f6, 0x434c6530,
                        0x2c5bad8f, 0x898881dc, 0x5f1c86d9, 0xc1f8e7f4));
    }

    #[test]
    fn test_rng_clone() {
        let seed : &[_] = &[0u32; 8];
        let mut rng: ChaChaRng = SeedableRng::from_seed(seed);
        let mut clone = rng.clone();
        for _ in 0..16 {
            assert_eq!(rng.next_u64(), clone.next_u64());
        }
    }
}
<commit_msg>Adding examples to ChaChaRng<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The ChaCha random number generator.

use std::num::Wrapping as w;
use {Rng, SeedableRng, Rand, w32};

const KEY_WORDS    : usize =  8; // 8 words for the 256-bit key
const STATE_WORDS  : usize = 16;
const CHACHA_ROUNDS: u32 = 20; // Cryptographically secure from 8 upwards as of this writing

/// A random number generator that uses the ChaCha20 algorithm [1].
///
/// The ChaCha algorithm is widely accepted as suitable for
/// cryptographic purposes, but this implementation has not been
/// verified as such. Prefer a generator like `OsRng` that defers to
/// the operating system for cases that need high security.
///
/// [1]: D. J. Bernstein, [*ChaCha, a variant of
/// Salsa20*](http://cr.yp.to/chacha.html)
#[derive(Copy, Clone)]
pub struct ChaChaRng {
    buffer:  [w32; STATE_WORDS], // Internal buffer of output
    state:   [w32; STATE_WORDS], // Initial state
    index:   usize,                 // Index into state
}

static EMPTY: ChaChaRng = ChaChaRng {
    buffer:  [w(0); STATE_WORDS],
    state:   [w(0); STATE_WORDS],
    index:   STATE_WORDS
};

macro_rules! quarter_round{
    ($a: expr, $b: expr, $c: expr, $d: expr) => {{
        $a = $a + $b; $d = $d ^ $a; $d = w($d.0.rotate_left(16));
        $c = $c + $d; $b = $b ^ $c; $b = w($b.0.rotate_left(12));
        $a = $a + $b; $d = $d ^ $a; $d = w($d.0.rotate_left( 8));
        $c = $c + $d; $b = $b ^ $c; $b = w($b.0.rotate_left( 7));
    }}
}

macro_rules! double_round{
    ($x: expr) => {{
        // Column round
        quarter_round!($x[ 0], $x[ 4], $x[ 8], $x[12]);
        quarter_round!($x[ 1], $x[ 5], $x[ 9], $x[13]);
        quarter_round!($x[ 2], $x[ 6], $x[10], $x[14]);
        quarter_round!($x[ 3], $x[ 7], $x[11], $x[15]);
        // Diagonal round
        quarter_round!($x[ 0], $x[ 5], $x[10], $x[15]);
        quarter_round!($x[ 1], $x[ 6], $x[11], $x[12]);
        quarter_round!($x[ 2], $x[ 7], $x[ 8], $x[13]);
        quarter_round!($x[ 3], $x[ 4], $x[ 9], $x[14]);
    }}
}

#[inline]
fn core(output: &mut [w32; STATE_WORDS], input: &[w32; STATE_WORDS]) {
    *output = *input;

    for _ in 0..CHACHA_ROUNDS / 2 {
        double_round!(output);
    }

    for i in 0..STATE_WORDS {
        output[i] = output[i] + input[i];
    }
}

impl ChaChaRng {

    /// Create an ChaCha random number generator using the default
    /// fixed key of 8 zero words.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use rand::{Rng, ChaChaRng};
    ///
    /// let mut ra = ChaChaRng::new_unseeded();
    /// println!("{:?}", ra.next_u32());
    /// println!("{:?}", ra.next_u32());
    /// ```
    ///
    /// Since this equivalent to a RNG with a fixed seed, repeated executions
    /// of an unseeded RNG will produce the same result. This code sample will
    /// consistently produce:
    /// - 2917185654
    /// - 2419978656
    pub fn new_unseeded() -> ChaChaRng {
        let mut rng = EMPTY;
        rng.init(&[0; KEY_WORDS]);
        rng
    }

    /// Sets the internal 128-bit ChaCha counter to
    /// a user-provided value. This permits jumping
    /// arbitrarily ahead (or backwards) in the pseudorandom stream.
    ///
    /// Since the nonce words are used to extend the counter to 128 bits,
    /// users wishing to obtain the conventional ChaCha pseudorandom stream
    /// associated with a particular nonce can call this function with
    /// arguments `0, desired_nonce`.
    ///
    /// # Examples
    ///
    /// ```
    /// use rand::{Rng, ChaChaRng};
    ///
    /// let mut ra = ChaChaRng::new_unseeded();
    /// ra.set_counter(0u64, 1234567890u64);
    /// println!("{:?}", ra.next_u32());
    /// println!("{:?}", ra.next_u32());
    /// ```
    pub fn set_counter(&mut self, counter_low: u64, counter_high: u64) {
        self.state[12] = w((counter_low >>  0) as u32);
        self.state[13] = w((counter_low >> 32) as u32);
        self.state[14] = w((counter_high >>  0) as u32);
        self.state[15] = w((counter_high >> 32) as u32);
        self.index = STATE_WORDS; // force recomputation
    }

    /// Initializes `self.state` with the appropriate key and constants
    ///
    /// We deviate slightly from the ChaCha specification regarding
    /// the nonce, which is used to extend the counter to 128 bits.
    /// This is provably as strong as the original cipher, though,
    /// since any distinguishing attack on our variant also works
    /// against ChaCha with a chosen-nonce. See the XSalsa20 [1]
    /// security proof for a more involved example of this.
    ///
    /// The modified word layout is:
    /// ```text
    /// constant constant constant constant
    /// key      key      key      key
    /// key      key      key      key
    /// counter  counter  counter  counter
    /// ```
    /// [1]: Daniel J. Bernstein. [*Extending the Salsa20
    /// nonce.*](http://cr.yp.to/papers.html#xsalsa)
    fn init(&mut self, key: &[u32; KEY_WORDS]) {
        self.state[0] = w(0x61707865);
        self.state[1] = w(0x3320646E);
        self.state[2] = w(0x79622D32);
        self.state[3] = w(0x6B206574);

        for i in 0..KEY_WORDS {
            self.state[4+i] = w(key[i]);
        }

        self.state[12] = w(0);
        self.state[13] = w(0);
        self.state[14] = w(0);
        self.state[15] = w(0);

        self.index = STATE_WORDS;
    }

    /// Refill the internal output buffer (`self.buffer`)
    fn update(&mut self) {
        core(&mut self.buffer, &self.state);
        self.index = 0;
        // update 128-bit counter
        self.state[12] = self.state[12] + w(1);
        if self.state[12] != w(0) { return };
        self.state[13] = self.state[13] + w(1);
        if self.state[13] != w(0) { return };
        self.state[14] = self.state[14] + w(1);
        if self.state[14] != w(0) { return };
        self.state[15] = self.state[15] + w(1);
    }
}

impl Rng for ChaChaRng {
    #[inline]
    fn next_u32(&mut self) -> u32 {
        if self.index == STATE_WORDS {
            self.update();
        }

        let value = self.buffer[self.index % STATE_WORDS];
        self.index += 1;
        value.0
    }
}

impl<'a> SeedableRng<&'a [u32]> for ChaChaRng {

    fn reseed(&mut self, seed: &'a [u32]) {
        // reset state
        self.init(&[0u32; KEY_WORDS]);
        // set key in place
        let key = &mut self.state[4 .. 4+KEY_WORDS];
        for (k, s) in key.iter_mut().zip(seed.iter()) {
            *k = w(*s);
        }
    }

    /// Create a ChaCha generator from a seed,
    /// obtained from a variable-length u32 array.
    /// Only up to 8 words are used; if less than 8
    /// words are used, the remaining are set to zero.
    fn from_seed(seed: &'a [u32]) -> ChaChaRng {
        let mut rng = EMPTY;
        rng.reseed(seed);
        rng
    }
}

impl Rand for ChaChaRng {
    fn rand<R: Rng>(other: &mut R) -> ChaChaRng {
        let mut key : [u32; KEY_WORDS] = [0; KEY_WORDS];
        for word in key.iter_mut() {
            *word = other.gen();
        }
        SeedableRng::from_seed(&key[..])
    }
}

#[cfg(test)]
mod test {
    use {Rng, SeedableRng};
    use super::ChaChaRng;

    #[test]
    fn test_rng_rand_seeded() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut ra: ChaChaRng = SeedableRng::from_seed(&s[..]);
        let mut rb: ChaChaRng = SeedableRng::from_seed(&s[..]);
        assert!(::test::iter_eq(ra.gen_ascii_chars().take(100),
                                rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_seeded() {
        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);
        let mut rb: ChaChaRng = SeedableRng::from_seed(seed);
        assert!(::test::iter_eq(ra.gen_ascii_chars().take(100),
                                rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_reseed() {
        let s = ::test::rng().gen_iter::<u32>().take(8).collect::<Vec<u32>>();
        let mut r: ChaChaRng = SeedableRng::from_seed(&s[..]);
        let string1: String = r.gen_ascii_chars().take(100).collect();

        r.reseed(&s);

        let string2: String = r.gen_ascii_chars().take(100).collect();
        assert_eq!(string1, string2);
    }

    #[test]
    fn test_rng_true_values() {
        // Test vectors 1 and 2 from
        // http://tools.ietf.org/html/draft-nir-cfrg-chacha20-poly1305-04
        let seed : &[_] = &[0u32; 8];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        let v = (0..16).map(|_| ra.next_u32()).collect::<Vec<_>>();
        assert_eq!(v,
                   vec!(0xade0b876, 0x903df1a0, 0xe56a5d40, 0x28bd8653,
                        0xb819d2bd, 0x1aed8da0, 0xccef36a8, 0xc70d778b,
                        0x7c5941da, 0x8d485751, 0x3fe02477, 0x374ad8b8,
                        0xf4b8436a, 0x1ca11815, 0x69b687c3, 0x8665eeb2));

        let v = (0..16).map(|_| ra.next_u32()).collect::<Vec<_>>();
        assert_eq!(v,
                   vec!(0xbee7079f, 0x7a385155, 0x7c97ba98, 0x0d082d73,
                        0xa0290fcb, 0x6965e348, 0x3e53c612, 0xed7aee32,
                        0x7621b729, 0x434ee69c, 0xb03371d5, 0xd539d874,
                        0x281fed31, 0x45fb0a51, 0x1f0ae1ac, 0x6f4d794b));

        let seed : &[_] = &[0,1,2,3,4,5,6,7];
        let mut ra: ChaChaRng = SeedableRng::from_seed(seed);

        // Store the 17*i-th 32-bit word,
        // i.e., the i-th word of the i-th 16-word block
        let mut v : Vec<u32> = Vec::new();
        for _ in 0..16 {
            v.push(ra.next_u32());
            for _ in 0..16 {
                ra.next_u32();
            }
        }

        assert_eq!(v,
                   vec!(0xf225c81a, 0x6ab1be57, 0x04d42951, 0x70858036,
                        0x49884684, 0x64efec72, 0x4be2d186, 0x3615b384,
                        0x11cfa18e, 0xd3c50049, 0x75c775f6, 0x434c6530,
                        0x2c5bad8f, 0x898881dc, 0x5f1c86d9, 0xc1f8e7f4));
    }

    #[test]
    fn test_rng_clone() {
        let seed : &[_] = &[0u32; 8];
        let mut rng: ChaChaRng = SeedableRng::from_seed(seed);
        let mut clone = rng.clone();
        for _ in 0..16 {
            assert_eq!(rng.next_u64(), clone.next_u64());
        }
    }
}
<|endoftext|>