//
// Copyright (c) 2017 anygone contributors (see AUTHORS file)
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice, this
//   list of conditions and the following disclaimer.
//
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// * Neither the name of anygone nor the names of its
//   contributors may be used to endorse or promote products derived from
//   this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//

package satp

import (
	"bytes"
	"math/rand"
	"testing"
	"time"
)

const (
	NUM_RANDOM_DATASETS = 1000
)

var (
	IPv4Type     = uint16(0x0800)
	IPv4PingData = [...]byte{
		0x45, 0x00, 0x00, 0x54, 0xB7, 0xCC, 0x40, 0x00, 0x40, 0x01, 0xAD, 0x73, 0xC0, 0xA8, 0x2A, 0x17,
		0xC0, 0xA8, 0x2A, 0x01, 0x08, 0x00, 0xB4, 0x04, 0x3D, 0xE2, 0x00, 0x01, 0xE6, 0x00, 0xB1, 0x59,
		0x00, 0x00, 0x00, 0x00, 0xA4, 0xEA, 0x0B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x11, 0x12, 0x13,
		0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F, 0x20, 0x21, 0x22, 0x23,
		0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F, 0x30, 0x31, 0x32, 0x33,
		0x34, 0x35, 0x36, 0x37}

	IPv6Type     = uint16(0x86DD)
	IPv6PingData = [...]byte{
		0x60, 0x0E, 0x9B, 0x1F, 0x00, 0x40, 0x3A, 0x40, 0x2A, 0x02, 0x12, 0x34, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x2A, 0x02, 0x12, 0x34, 0x00, 0x00, 0x00, 0x00,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x42, 0x80, 0x00, 0xE6, 0x21, 0x46, 0x22, 0x00, 0x01,
		0x3E, 0x06, 0xB1, 0x59, 0x00, 0x00, 0x00, 0x00, 0x21, 0x5E, 0x0B, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F,
		0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F,
		0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37}

	EthernetType   = uint16(0x6558)
	ARPRequestData = [...]byte{
		0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xC6, 0x5A, 0xD4, 0xC7, 0x23, 0x43, 0x08, 0x06, 0x00, 0x01,
		0x08, 0x00, 0x06, 0x04, 0x00, 0x01, 0xC6, 0x5A, 0xD4, 0xC7, 0x23, 0x43, 0xC0, 0xA8, 0x2A, 0x17,
		0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xC0, 0xA8, 0x2A, 0x01}

	tRand *rand.Rand
)

func init() {
	tRand = rand.New(rand.NewSource(time.Now().UnixNano()))
}

func TestPlainPacketHeader(t *testing.T) {
	testvectors := []uint16{0, IPv4Type, IPv6Type, EthernetType, 0xAA55, 0xF00F}

	for _, vector := range testvectors {
		pkt := NewPlainPacket()
		if vector > 0 {
			pkt.SetPayloadType(vector)
		}
		result := pkt.getPacket()
		expected := []byte{byte(vector >> 8), byte(vector)}
		if bytes.Compare(result, expected) != 0 {
			t.Fatalf("resulting packet is invalid, is: '%v', should be '%v'", result, expected)
		}

		payloadType := pkt.GetPayloadType()
		if vector != payloadType {
			t.Fatalf("resulting payload type is invalid, is: 0x%04X, should be 0x%04X", payloadType, vector)
		}
	}
}

func TestPlainPacketReadFrom(t *testing.T) {
	testvectors := []struct {
		Data []byte
		Type uint16
	}{
		{IPv4PingData[:], IPv4Type},
		{IPv6PingData[:], IPv6Type},
		{ARPRequestData[:], EthernetType},
	}

	for _, vector := range testvectors {
		pkt := NewPlainPacket()
		pkt.SetPayloadType(vector.Type)
		n, err := pkt.ReadFrom(bytes.NewReader(vector.Data))
		if err != nil {
			t.Fatal("unexpected error:", err)
		}
		if n != int64(len(vector.Data)) {
			t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(vector.Data))
		}
		expected := []byte{}
		expected = append(expected, []byte{byte(vector.Type >> 8), byte(vector.Type)}...)
		expected = append(expected, vector.Data...)

		out := pkt.getPacket()
		if bytes.Compare(out, expected) != 0 {
			t.Fatalf("resulting packet is invalid, is: '%v', should be '%v'", out, expected)
		}
	}
}

func TestPlainPacketWriteTo(t *testing.T) {
	testvectors := []struct {
		Data []byte
		Type uint16
	}{
		{IPv4PingData[:], IPv4Type},
		{IPv6PingData[:], IPv6Type},
		{ARPRequestData[:], EthernetType},
	}

	for _, vector := range testvectors {
		pkt := NewPlainPacket()
		in := pkt.buffer[:0]
		in = append(in, []byte{byte(vector.Type >> 8), byte(vector.Type)}...)
		in = append(in, vector.Data...)
		pkt.payload = pkt.payload[:len(vector.Data)]

		out := &bytes.Buffer{}
		n, err := pkt.WriteTo(out)
		if err != nil {
			t.Fatal("unexpected error:", err)
		}
		if n != int64(len(vector.Data)) {
			t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(vector.Data))
		}

		outType := pkt.GetPayloadType()
		if outType != vector.Type {
			t.Fatalf("returned payload type is invalid, is: 0x%04X, should be 0x%04X", outType, vector.Type)
		}
		outData := out.Bytes()
		if bytes.Compare(outData, vector.Data) != 0 {
			t.Fatalf("resulting packet is invalid, is: '%v', should be '%v'", outData, vector.Data)
		}
	}
}

func generateRandomTestDataPlainPacket() (payloadType uint16, payload []byte) {
	payloadType = uint16(tRand.Uint32())

	packetLen := uint(tRand.Int31n(PACKET_BUFFER_SIZE - 2))
	payload = make([]byte, packetLen)
	tRand.Read(payload)
	return
}

func TestPlainPacketReadWrite(t *testing.T) {
	for i := 0; i < NUM_RANDOM_DATASETS; i++ {
		pkt := NewPlainPacket()
		inType, in := generateRandomTestDataPlainPacket()

		pkt.SetPayloadType(inType)
		n, err := pkt.ReadFrom(bytes.NewReader(in))
		if err != nil {
			t.Fatal("unexpected error:", err)
		}
		if n != int64(len(in)) {
			t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(in))
		}

		out := &bytes.Buffer{}
		if n, err = pkt.WriteTo(out); err != nil {
			t.Fatal("unexpected error:", err)
		}
		if n != int64(len(in)) {
			t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(in))
		}

		if bytes.Compare(in, out.Bytes()) != 0 {
			t.Fatalf("resulting packet is invalid, is: '%v', should be '%v'", out.Bytes(), in)
		}
	}
}

func TestEncryptedPacketHeader(t *testing.T) {
	testvectors := []struct {
		SequenceNumber uint32
		SenderID       uint16
		Mux            uint16
	}{
		{0, 0, 0},
		{23, 0, 0},
		{0, 12, 0},
		{0, 0, 42},
		{1, 2, 3},
		{4294967295, 0, 0},
		{0, 65535, 0},
		{0, 0, 65535},
		{1254234, 2, 23},
	}

	for _, vector := range testvectors {
		pkt := NewEncryptedPacket()
		if vector.SequenceNumber > 0 {
			pkt.SetSequenceNumber(vector.SequenceNumber)
		}
		if vector.SenderID > 0 {
			pkt.SetSenderID(vector.SenderID)
		}
		if vector.Mux > 0 {
			pkt.SetMux(vector.Mux)
		}
		result := pkt.getPacket()
		expected := []byte{
			byte(vector.SequenceNumber >> 24), byte(vector.SequenceNumber >> 16),
			byte(vector.SequenceNumber >> 8), byte(vector.SequenceNumber),
			byte(vector.SenderID >> 8), byte(vector.SenderID),
			byte(vector.Mux >> 8), byte(vector.Mux)}

		if bytes.Compare(result, expected) != 0 {
			t.Fatalf("resulting packet is invalid, is: '%v', should be '%v'", result, expected)
		}

		sequenceNumber := pkt.GetSequenceNumber()
		if vector.SequenceNumber != sequenceNumber {
			t.Fatalf("resulting sequence number is invalid, is: 0x%08X, should be 0x%08X", sequenceNumber, vector.SequenceNumber)
		}
		senderID := pkt.GetSenderID()
		if vector.SenderID != senderID {
			t.Fatalf("resulting sequence number is invalid, is: 0x%04X, should be 0x%04X", senderID, vector.SenderID)
		}
		Mux := pkt.GetMux()
		if vector.Mux != Mux {
			t.Fatalf("resulting sequence number is invalid, is: 0x%04X, should be 0x%04X", Mux, vector.Mux)
		}
	}
}

func TestEncryptedPacketSetAuthTagLen(t *testing.T) {
	var buffer [PACKET_BUFFER_SIZE]byte
	tRand.Read(buffer[:])

	testvectors := []struct {
		Data       []byte
		AuthTagLen int
		valid      bool
	}{
		{buffer[:10], 0, true},
		{buffer[:10], 1, false},
		{buffer[:11], 1, true},
		{buffer[:200], 20, true},
		{buffer[:200], 190, true},
		{buffer[:200], 191, false},
		{buffer[:1440], 20, true},
	}

	for _, vector := range testvectors {
		pkt := NewEncryptedPacket()
		n, err := pkt.ReadFrom(bytes.NewReader(vector.Data))
		if err != nil {
			t.Fatal("unexpected error:", err)
		}

		err = pkt.SetAuthTagLength(vector.AuthTagLen)
		if vector.valid {
			if err != nil {
				t.Fatal("unexpected error:", err)
			}
			if n != int64(len(vector.Data)) {
				t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(vector.Data))
			}

			out := pkt.GetAuthTagLength()
			if out != vector.AuthTagLen {
				t.Fatalf("resulting auth-tag length is wrong, is: %d, should be %d", out, vector.AuthTagLen)
			}
		} else {
			if err == nil {
				t.Fatalf("setting auth-tag length to %d on packet with length %d should give an error", vector.AuthTagLen, len(vector.Data))
			}
		}
	}

}

func TestEncryptedPacketReadFrom(t *testing.T) {
	testvectors := []struct {
		Data           []byte
		SequenceNumber uint32
		SenderID       uint16
		Mux            uint16
		valid          bool
	}{
		{[]byte{}, 0, 0, 0, false},
		{[]byte{0, 0, 0}, 0, 0, 0, false},
		{[]byte{0, 0, 0, 0, 0, 0, 0}, 0, 0, 0, false},
		{[]byte{0, 0, 0, 0, 0, 0, 0, 0, 0}, 0, 0, 0, false},
		{[]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, 0, 0, 0, true},
		{[]byte{1, 2, 3, 4, 5, 6, 7, 8, 0, 0}, 0x01020304, 0x0506, 0x0708, true},
	}

	for _, vector := range testvectors {
		pkt := NewEncryptedPacket()
		n, err := pkt.ReadFrom(bytes.NewReader(vector.Data))
		if vector.valid {
			if err != nil {
				t.Fatal("unexpected error:", err)
			}
			if n != int64(len(vector.Data)) {
				t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(vector.Data))
			}
			seq := pkt.GetSequenceNumber()
			if seq != vector.SequenceNumber {
				t.Fatalf("returned sequence number, is: 0x%08X, should be 0x%08X", seq, vector.SequenceNumber)
			}
			sid := pkt.GetSenderID()
			if sid != vector.SenderID {
				t.Fatalf("returned sender id, is: 0x%04X, should be 0x%04X", sid, vector.SenderID)
			}
			mux := pkt.GetMux()
			if mux != vector.Mux {
				t.Fatalf("returned mux, is: 0x%04X, should be 0x%04X", mux, vector.Mux)
			}

			expected := vector.Data[8:]
			out := pkt.payload
			if bytes.Compare(out, expected) != 0 {
				t.Fatalf("resulting payload is invalid, is: '%v', should be '%v'", out, expected)
			}
		} else {
			if err == nil {
				t.Fatalf("reading packet from %v should give an error", vector.Data)
			}
		}
	}
}

func TestEncryptedPacketWriteTo(t *testing.T) {
	var buffer [PACKET_BUFFER_SIZE]byte
	tRand.Read(buffer[:])

	testvectors := []struct {
		Data           []byte
		SequenceNumber uint32
		SenderID       uint16
		Mux            uint16
	}{
		{[]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, 0, 0, 0},
		{[]byte{1, 2, 3, 4, 5, 6, 7, 8, 17, 32}, 0x01020304, 0x0506, 0x0708},
	}

	for _, vector := range testvectors {
		pkt := NewEncryptedPacket()
		pkt.SetSequenceNumber(vector.SequenceNumber)
		pkt.SetSenderID(vector.SenderID)
		pkt.SetMux(vector.Mux)
		pkt.payload = pkt.payload[:len(vector.Data)-8]
		copy(pkt.payload, vector.Data[8:])

		w := &bytes.Buffer{}
		n, err := pkt.WriteTo(w)
		if err != nil {
			t.Fatal("unexpected error:", err)
		}
		if n != int64(len(vector.Data)) {
			t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(vector.Data))
		}
		out := w.Bytes()
		if bytes.Compare(out, vector.Data) != 0 {
			t.Fatalf("resulting payload is invalid, is: '%v', should be '%v'", out, vector.Data)
		}
	}
}

func generateRandomTestDataEncryptedPacket() (authTagLen int, packetData []byte) {
	minLen := int32(8 + 2)
	packetLen := uint(minLen + tRand.Int31n(PACKET_BUFFER_SIZE-minLen))
	if packetLen > uint(minLen) {
		authTagLen = int(tRand.Int31n(int32(packetLen) - minLen))
	}
	packetData = make([]byte, packetLen)
	tRand.Read(packetData)
	return
}

func TestEncryptedPacketReadWrite(t *testing.T) {
	for i := 0; i < NUM_RANDOM_DATASETS; i++ {
		pkt := NewEncryptedPacket()
		authTagLen, in := generateRandomTestDataEncryptedPacket()

		n, err := pkt.ReadFrom(bytes.NewReader(in))
		if err != nil {
			t.Fatal("unexpected error:", err)
		}
		if n != int64(len(in)) {
			t.Fatalf("returned length is invalid, is: %d, should be %d", n, len(in))
		}

		err = pkt.SetAuthTagLength(authTagLen)
		if err != nil {
			t.Fatal("unexpected error:", err)
		}

		out := &bytes.Buffer{}
		if n, err = pkt.WriteTo(out); err != nil {
			t.Fatal("unexpected error:", err)
		}
		if n != int64(len(in)) {
			t.Fatalf("returned length is invalid, is: %d, should be %d ... authTagLength is %d", n, len(in), authTagLen)
		}

		if bytes.Compare(in, out.Bytes()) != 0 {
			t.Fatalf("resulting packet is invalid, is: '%v', should be '%v'", out.Bytes(), in)
		}
	}
}

//
// Benchmarking
//

type devZero struct {
}

func (dz *devZero) Read(data []byte) (int, error) {
	return len(data), nil
}

func (dz *devZero) Write(data []byte) (int, error) {
	return len(data), nil
}

func BenchmarkPlainPacketReadFrom(b *testing.B) {
	in := &devZero{}
	pkt := NewPlainPacket()

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if _, err := pkt.ReadFrom(in); err != nil {
			b.Fatal("unexpected error:", err)
		}
	}
}

func BenchmarkPlainPacketWriteTo(b *testing.B) {
	ins := []*PlainPacket{}
	for i := 0; i < NUM_RANDOM_DATASETS; i++ {
		_, payload := generateRandomTestDataPlainPacket()
		in := NewPlainPacket()
		if _, err := in.ReadFrom(bytes.NewReader(payload)); err != nil {
			b.Fatal("unexpected error:", err)
		}
		ins = append(ins, in)
	}
	out := &devZero{}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		in := ins[i%len(ins)]
		if _, err := in.WriteTo(out); err != nil {
			b.Fatal("unexpected error:", err)
		}
	}
}

func BenchmarkEncryptedPacketReadFrom(b *testing.B) {
	in := &devZero{}
	pkt := NewEncryptedPacket()

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if _, err := pkt.ReadFrom(in); err != nil {
			b.Fatal("unexpected error:", err)
		}
	}
}

func BenchmarkEncryptedPacketWriteTo(b *testing.B) {
	ins := []*EncryptedPacket{}
	for i := 0; i < NUM_RANDOM_DATASETS; i++ {
		_, data := generateRandomTestDataEncryptedPacket()
		in := NewEncryptedPacket()
		if _, err := in.ReadFrom(bytes.NewReader(data)); err != nil {
			b.Fatal("unexpected error:", err)
		}
		ins = append(ins, in)
	}
	out := &devZero{}

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		in := ins[i%len(ins)]
		if _, err := in.WriteTo(out); err != nil {
			b.Fatal("unexpected error:", err)
		}
	}
}