// Copyright 2015-2016 Brian Smith.
// Portions Copyright (c) 2014, 2015, Google Inc.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHORS DISCLAIM ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY
// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

// TODO: enforce maximum input length.

// Work around compiler bug?
#![allow(non_shorthand_field_patterns)]

use {c, chacha, constant_time, error, polyfill};
use core;

// The assembly functions we call expect the state to be 8-byte aligned. We do
// this manually, by copying the data into an aligned slice, rather than by
// asking the compiler to align the `Opaque` struct.
fn with_aligned<F>(opaque: &mut Opaque, f: F)
    where F: FnOnce(&mut Opaque)
{
    let mut buf = [0u8; OPAQUE_LEN + 7];
    let aligned_start = (buf.as_ptr() as usize + 7) & !7;
    let offset = aligned_start - (buf.as_ptr() as usize);
    let aligned_buf =
        slice_as_array_ref_mut!(
            &mut buf[offset..(offset + OPAQUE_LEN)], OPAQUE_LEN).unwrap();
    aligned_buf.copy_from_slice(&opaque[..]);
    f(aligned_buf);
    opaque.copy_from_slice(aligned_buf);
}

impl SigningContext {
    #[inline]
    pub fn from_key(key: Key) -> SigningContext {
        #[inline]
        fn read_u32(buf: &[u8]) -> u32 {
            polyfill::slice::u32_from_le_u8(slice_as_array_ref!(buf, 4).unwrap())
        }

        let (key, nonce) = key.bytes.split_at(16);
        let key = slice_as_array_ref!(key, 16).unwrap();

        let mut ctx = SigningContext {
            opaque: [0u8; OPAQUE_LEN],
            // TODO: When we can get explicit alignment, make `nonce` an
            // aligned `u8[16]` and get rid of this `u8[16]` -> `u32[4]`
            // conversion.
            nonce: [
                read_u32(&nonce[0..4]),
                read_u32(&nonce[4..8]),
                read_u32(&nonce[8..12]),
                read_u32(&nonce[12..16]),
            ],
            buf: [0; BLOCK_LEN],
            buf_used: 0,
            func: Funcs {
                blocks_fn: GFp_poly1305_blocks,
                emit_fn: GFp_poly1305_emit
            }
        };

        { // borrow ctx
            let &mut SigningContext {
                opaque: ref mut opaque,
                func: ref mut func,
                ..
            } = &mut ctx;
            with_aligned(opaque, |opaque| {
            // On some platforms `init()` doesn't initialize `funcs`. The
            // return value of `init()` indicates whether it did or not. Since
            // we already gave `func` a default value above, we can ignore the
            // return value assuming `init()` doesn't change `func` if it chose
            // not to initialize it. Note that this is different than what
            // BoringSSL does.
                let _ = init(opaque, key, func);
            });
        }

        ctx
    }

    pub fn update(&mut self, mut input: &[u8]) {
        let &mut SigningContext {
            opaque: ref mut opaque,
            buf: ref mut buf,
            buf_used: ref mut buf_used,
            func: ref func,
            ..
        } = self;
        with_aligned(opaque, |opaque: &mut Opaque| {
            if *buf_used != 0 {
                let todo = core::cmp::min(input.len(), BLOCK_LEN - *buf_used);

                buf[*buf_used..(*buf_used + todo)].copy_from_slice(
                    &input[..todo]);
                *buf_used += todo;
                input = &input[todo..];

                if *buf_used == BLOCK_LEN {
                    func.blocks(opaque, buf, Pad::Pad);
                    *buf_used = 0;
                }
            }

            if input.len() >= BLOCK_LEN {
                let todo = input.len() & !(BLOCK_LEN - 1);
                let (complete_blocks, remainder) = input.split_at(todo);
                func.blocks(opaque, complete_blocks, Pad::Pad);
                input = remainder;
            }

            if input.len() != 0 {
                buf[..input.len()].copy_from_slice(input);
                *buf_used = input.len();
            }
        });
    }

    pub fn sign(mut self, tag_out: &mut Tag) {
        let &mut SigningContext {
            opaque: ref mut opaque,
            nonce: ref nonce,
            buf: ref mut buf,
            buf_used: buf_used,
            func: ref func,
        } = &mut self;
        with_aligned(opaque, |opaque| {
            if buf_used != 0 {
                buf[buf_used] = 1;
                for byte in &mut buf[(buf_used + 1)..] {
                    *byte = 0;
                }
                func.blocks(opaque, &buf[..], Pad::AlreadyPadded);
            }

            func.emit(opaque, tag_out, nonce);
        });
    }
}

pub fn verify(key: Key, msg: &[u8], tag: &Tag)
              -> Result<(), error::Unspecified> {
    let mut calculated_tag = [0u8; TAG_LEN];
    sign(key, msg, &mut calculated_tag);
    constant_time::verify_slices_are_equal(&calculated_tag[..], tag)
}

pub fn sign(key: Key, msg: &[u8], tag: &mut Tag) {
    let mut ctx = SigningContext::from_key(key);
    ctx.update(msg);
    ctx.sign(tag)
}

#[cfg(test)]
pub fn check_state_layout() {
    let required_state_size =
        if cfg!(target_arch = "x86") {
            // See comment above `_poly1305_init_sse2` in poly1305-x86.pl.
            Some(4 * (5 + 1 + 4 + 2 + 4 * 9))
        } else if cfg!(target_arch = "x86_64") {
            // See comment above `__poly1305_block` in poly1305-x86_64.pl.
            Some(4 * (5 + 1 + 2 * 2 + 2 + 4 * 9))
        } else {
            // TODO(davidben): Figure out the layout of the struct. For now,
            // `OPAQUE_LEN` is taken from OpenSSL.
            None
        };

    if let Some(required_state_size) = required_state_size {
        assert!(core::mem::size_of::<Opaque>() >= required_state_size);
    }
}

/// A Poly1305 key.
pub struct Key {
    bytes: KeyAndNonceBytes,
}

impl Key {
    pub fn derive_using_chacha(chacha20_key: &chacha::Key,
                               counter: &chacha::Counter) -> Key {
        let mut bytes = [0u8; KEY_LEN];
        chacha::chacha20_xor_in_place(chacha20_key, counter, &mut bytes);
        Key { bytes }
    }

    #[cfg(test)]
    pub fn from_test_vector(bytes: &[u8; KEY_LEN]) -> Key {
        Key { bytes: *bytes }
    }
}

type KeyAndNonceBytes = [u8; 2 * BLOCK_LEN];

type KeyBytes = [u8; BLOCK_LEN];
type Nonce = [u32; BLOCK_LEN / 4];

/// The length of a `key`.
pub const KEY_LEN: usize = 32;

/// A Poly1305 tag.
pub type Tag = [u8; TAG_LEN];

/// The length of a `Tag`.
pub const TAG_LEN: usize = BLOCK_LEN;

const BLOCK_LEN: usize = 16;

/// The memory manipulated by the assembly.
type Opaque = [u8; OPAQUE_LEN];

const OPAQUE_LEN: usize = 192;

#[repr(C)]
struct Funcs {
    blocks_fn: unsafe extern fn(&mut Opaque, input: *const u8,
                                input_len: c::size_t, should_pad: Pad),
    emit_fn: unsafe extern fn(&mut Opaque, &mut Tag, nonce: &Nonce),
}

#[inline]
fn init(state: &mut Opaque, key: &KeyBytes, func: &mut Funcs) -> i32 {
    debug_assert_eq!(state.as_ptr() as usize % 8, 0);
    unsafe {
        GFp_poly1305_init_asm(state, key, func)
    }
}

#[repr(u32)]
enum Pad {
    AlreadyPadded = 0,
    Pad = 1,
}

impl Funcs {
    #[inline]
    fn blocks(&self, state: &mut Opaque, data: &[u8], should_pad: Pad) {
        debug_assert_eq!(state.as_ptr() as usize % 8, 0);
        unsafe {
            (self.blocks_fn)(state, data.as_ptr(), data.len(), should_pad);
        }
    }

    #[inline]
    fn emit(&self, state: &mut Opaque, tag_out: &mut Tag, nonce: &Nonce) {
        debug_assert_eq!(state.as_ptr() as usize % 8, 0);
        unsafe {
             (self.emit_fn)(state, tag_out, nonce);
        }
    }
}

pub struct SigningContext {
    opaque: Opaque,
    nonce: [u32; 4],
    buf: [u8; BLOCK_LEN],
    buf_used: usize,
    func: Funcs
}

extern {
    fn GFp_poly1305_init_asm(state: &mut Opaque, key: &KeyBytes,
                             out_func: &mut Funcs) -> c::int;
    fn GFp_poly1305_blocks(state: &mut Opaque, input: *const u8, len: c::size_t,
                           should_pad: Pad);
    fn GFp_poly1305_emit(state: &mut Opaque, mac: &mut Tag, nonce: &Nonce);
}

#[cfg(test)]
mod tests {
    use {error, test};
    use core;
    use super::*;

    #[test]
    pub fn test_state_layout() {
        check_state_layout();
    }

    // Adapted from BoringSSL's crypto/poly1305/poly1305_test.cc.
    #[test]
    pub fn test_poly1305() {
        test::from_file("src/poly1305_test.txt", |section, test_case| {
            assert_eq!(section, "");
            let key = test_case.consume_bytes("Key");
            let key = slice_as_array_ref!(&key, KEY_LEN).unwrap();
            let input = test_case.consume_bytes("Input");
            let expected_mac = test_case.consume_bytes("MAC");
            let expected_mac =
                slice_as_array_ref!(&expected_mac, TAG_LEN).unwrap();

            // Test single-shot operation.
            {
                let key = Key::from_test_vector(&key);
                let mut ctx = SigningContext::from_key(key);
                ctx.update(&input);
                let mut actual_mac = [0; TAG_LEN];
                ctx.sign(&mut actual_mac);
                assert_eq!(&expected_mac[..], &actual_mac[..]);
            }
            {
                let key = Key::from_test_vector(&key);
                let mut actual_mac = [0; TAG_LEN];
                sign(key, &input, &mut actual_mac);
                assert_eq!(&expected_mac[..], &actual_mac[..]);
            }
            {
                let key = Key::from_test_vector(&key);
                assert_eq!(Ok(()), verify(key, &input, &expected_mac));
            }

            // Test streaming byte-by-byte.
            {
                let key = Key::from_test_vector(&key);
                let mut ctx = SigningContext::from_key(key);
                for chunk in input.chunks(1) {
                    ctx.update(chunk);
                }
                let mut actual_mac = [0u8; TAG_LEN];
                ctx.sign(&mut actual_mac);
                assert_eq!(&expected_mac[..], &actual_mac[..]);
            }

            test_poly1305_simd(0, key, &input, expected_mac)?;
            test_poly1305_simd(16, key, &input, expected_mac)?;
            test_poly1305_simd(32, key, &input, expected_mac)?;
            test_poly1305_simd(48, key, &input, expected_mac)?;

            Ok(())
        })
    }

    fn test_poly1305_simd(excess: usize, key: &[u8; KEY_LEN], input: &[u8],
                          expected_mac: &[u8; TAG_LEN])
                          -> Result<(), error::Unspecified> {
        let key = Key::from_test_vector(&key);
        let mut ctx = SigningContext::from_key(key);

        // Some implementations begin in non-SIMD mode and upgrade on demand.
        // Stress the upgrade path.
        let init = core::cmp::min(input.len(), 16);
        ctx.update(&input[..init]);

        let long_chunk_len = 128 + excess;
        for chunk in input[init..].chunks(long_chunk_len + excess) {
            if chunk.len() > long_chunk_len {
                let (long, short) = chunk.split_at(long_chunk_len);

                // Feed 128 + |excess| bytes to test SIMD mode.
                ctx.update(long);

                // Feed |excess| bytes to ensure SIMD mode can handle short
                // inputs.
                ctx.update(short);
            } else {
                // Handle the last chunk.
                ctx.update(chunk);
            }
        }

        let mut actual_mac = [0u8; TAG_LEN];
        ctx.sign(&mut actual_mac);
        assert_eq!(&expected_mac[..], &actual_mac);

        Ok(())
    }
}
