package chacha20

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"testing"
)

// rfcKey and rfcNonce are the inputs from RFC 8439, Section 2.4.2.
var (
	rfcKey   = [32]byte{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}
	rfcNonce = [12]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00}
)

// TestBlockRFC8439 checks the block function against the state matrices and
// keystream published in RFC 8439, Section 2.4.2 (initial counter 1).
func TestBlockRFC8439(t *testing.T) {
	wantWords := [2][16]uint32{
		{
			0xf3514f22, 0xe1d91b40, 0x6f27de2f, 0xed1d63b8,
			0x821f138c, 0xe2062c3d, 0xecca4f7e, 0x78cff39e,
			0xa30a3b8a, 0x920a6072, 0xcd7479b5, 0x34932bed,
			0x40ba4c79, 0xcd343ec6, 0x4c2c21ea, 0xb7417df0,
		},
		{
			0x9f74a669, 0x410f633f, 0x28feca22, 0x7ec44dec,
			0x6d34d426, 0x738cb970, 0x3ac5e9f3, 0x45590cc4,
			0xda6e8b39, 0x892c831a, 0xcdea67c1, 0x2b7e1d90,
			0x037463f3, 0xa11a2073, 0xe8bcfb88, 0xedc49139,
		},
	}

	for i, counter := range []uint32{1, 2} {
		got := Block(&rfcKey, counter, &rfcNonce)
		for j, want := range wantWords[i] {
			if v := binary.LittleEndian.Uint32(got[4*j : 4*j+4]); v != want {
				t.Fatalf("counter %d word %d = %08x, want %08x", counter, j, v, want)
			}
		}
	}

	const keystreamHex = "" +
		"224f51f3401bd9e12fde276fb8631ded8c131f823d2c06e27e4fcaec9ef3cf78" +
		"8a3b0aa372600a92b57974cded2b9334794cba40c63e34cdea212c4cf07d41b7" +
		"69a6749f3f630f4122cafe28ec4dc47e26d4346d70b98c73f3e9c53ac40c5945" +
		"398b6eda1a832c89c167eacd901d7e2bf363"
	wantKeystream, err := hex.DecodeString(keystreamHex)
	if err != nil {
		t.Fatal(err)
	}
	if len(wantKeystream) != 114 {
		t.Fatalf("bad test vector length %d", len(wantKeystream))
	}

	// The RFC keystream spans two blocks starting at counter 1.
	var got []byte
	b1 := Block(&rfcKey, 1, &rfcNonce)
	b2 := Block(&rfcKey, 2, &rfcNonce)
	got = append(got, b1[:]...)
	got = append(got, b2[:]...)
	got = got[:len(wantKeystream)]
	if !bytes.Equal(got, wantKeystream) {
		t.Fatalf("keystream mismatch:\n got %x\nwant %x", got, wantKeystream)
	}
}

// TestXORKeyStreamRFC8439 encrypts the sunscreen plaintext from RFC 8439,
// Section 2.4.2 and compares against the published ciphertext.
func TestXORKeyStreamRFC8439(t *testing.T) {
	plaintext := []byte("Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it.")
	const ciphertextHex = "" +
		"6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0b" +
		"f91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d8" +
		"07ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab7793736" +
		"5af90bbf74a35be6b40b8eedf2785e42874d"
	wantCiphertext, err := hex.DecodeString(ciphertextHex)
	if err != nil {
		t.Fatal(err)
	}

	// XORKeyStream always starts at counter 0, so prepend a block's worth of
	// keystream: XORing the first 64 dummy bytes consumes counter 0, and the
	// RFC plaintext is then encrypted with counters 1 and 2.
	dummy := make([]byte, 64)
	buf := append(append([]byte{}, dummy...), plaintext...)
	XORKeyStream(&rfcKey, &rfcNonce, buf, buf)

	got := buf[64:]
	if !bytes.Equal(got, wantCiphertext) {
		t.Fatalf("ciphertext mismatch:\n got %x\nwant %x", got, wantCiphertext)
	}
}

// TestXORKeyStreamCounterStartsAtZero verifies that the first block uses
// counter 0 and that subsequent blocks increment the counter.
func TestXORKeyStreamCounterStartsAtZero(t *testing.T) {
	var key [32]byte
	var nonce [12]byte
	for i := range key {
		key[i] = byte(i)
	}
	for i := range nonce {
		nonce[i] = byte(0xa0 + i)
	}
	src := make([]byte, 3*BlockSize)
	for i := range src {
		src[i] = byte(i)
	}

	dst := make([]byte, len(src))
	XORKeyStream(&key, &nonce, dst, src)

	want := make([]byte, len(src))
	for b := 0; b < 3; b++ {
		ks := Block(&key, uint32(b), &nonce)
		for i := 0; i < BlockSize; i++ {
			want[b*BlockSize+i] = src[b*BlockSize+i] ^ ks[i]
		}
	}
	if !bytes.Equal(dst, want) {
		t.Fatal("keystream is not the counter 0,1,2 sequence")
	}
}

// TestXORKeyStreamAlias checks in-place operation and non-overlapping output.
func TestXORKeyStreamAlias(t *testing.T) {
	var key [32]byte
	var nonce [12]byte
	for i := range key {
		key[i] = byte(0x11 * i)
	}
	for i := range nonce {
		nonce[i] = byte(0x30 + i)
	}
	plaintext := make([]byte, 201) // not a multiple of the block size
	for i := range plaintext {
		plaintext[i] = byte(7*i + 3)
	}

	// dst and src are the same slice.
	aliased := append([]byte{}, plaintext...)
	XORKeyStream(&key, &nonce, aliased, aliased)
	if bytes.Equal(aliased, plaintext) {
		t.Fatal("in-place XOR did not change the buffer")
	}

	// dst and src are distinct.
	separate := make([]byte, len(plaintext))
	XORKeyStream(&key, &nonce, separate, plaintext)
	if !bytes.Equal(aliased, separate) {
		t.Fatal("aliased and non-aliased results differ")
	}
	ciphertext := append([]byte{}, separate...)

	// Round trip restores the plaintext in both cases.
	XORKeyStream(&key, &nonce, aliased, aliased)
	if !bytes.Equal(aliased, plaintext) {
		t.Fatal("in-place round trip failed")
	}
	XORKeyStream(&key, &nonce, separate, separate)
	if !bytes.Equal(separate, plaintext) {
		t.Fatal("round trip failed")
	}

	// A dst longer than src leaves the tail untouched.
	long := make([]byte, len(plaintext)+17)
	for i := range long {
		long[i] = 0xee
	}
	XORKeyStream(&key, &nonce, long, plaintext)
	if !bytes.Equal(long[len(plaintext):], bytes.Repeat([]byte{0xee}, 17)) {
		t.Fatal("wrote past len(src)")
	}
	if !bytes.Equal(long[:len(plaintext)], ciphertext) {
		t.Fatal("short-copy mismatch")
	}
}

// TestXORKeyStreamEmpty checks the zero-length case.
func TestXORKeyStreamEmpty(t *testing.T) {
	var key [32]byte
	var nonce [12]byte
	XORKeyStream(&key, &nonce, nil, nil)
	XORKeyStream(&key, &nonce, []byte{}, []byte{})
}

func TestXORKeyStreamShortDstPanics(t *testing.T) {
	var key [32]byte
	var nonce [12]byte
	defer func() {
		if recover() == nil {
			t.Fatal("expected panic for short destination")
		}
	}()
	XORKeyStream(&key, &nonce, make([]byte, 3), make([]byte, 4))
}