//
// 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 (
	"crypto/rand"
	"net"
	"testing"
)

func TestSecurityAssociationNew(t *testing.T) {
	testvectors := []struct {
		numEndpoints uint
		initSeqOut   uint32
		initSeqIn    uint32
	}{
		{0, 0, 0},
		{1, 0, 0},
		{3, 12, 0},
		{17, 0, 144},
		{17, 124, 144},
	}

	for _, vector := range testvectors {
		sa := NewSecurityAssociation(nil, vector.numEndpoints, vector.initSeqOut, vector.initSeqIn, 0)
		if sa == nil {
			t.Fatal("NewSecurityAssociation returned nil")
		}
		if sa.endpoints == nil {
			t.Fatal("endpoints must not be nil")
		}
		if uint(len(sa.endpoints)) != vector.numEndpoints {
			t.Fatalf("wrong number of endpoints is %d but should be %d", len(sa.endpoints), vector.numEndpoints)
		}
	}
}

func TestSecurityAssociationGenerate(t *testing.T) {
	var keymat [46]byte
	rand.Read(keymat[:])

	kd, err := NewAESCTRKeyDerivation(keymat[:32], keymat[32:], RoleLeft)
	if err != nil {
		t.Fatal("unexpected error:", err)
	}

	sa := NewSecurityAssociation(kd, 1, 0, 0, 0)

	var out [32]byte
	err = sa.KeyGenerate(Outbound, UsageEncryptKey, 23, out[:32])
	if err != nil {
		t.Fatal("unexpected error:", err)
	}
	err = sa.KeyGenerate(Outbound, UsageEncryptSalt, 23, out[:14])
	if err != nil {
		t.Fatal("unexpected error:", err)
	}
}

func TestSecurityAssociationEndpointsEqual(t *testing.T) {
	local4Num444, _ := net.ResolveUDPAddr("udp4", "127.0.0.1:444")
	local4Name444, _ := net.ResolveUDPAddr("udp4", "localhost:444")
	local4Num1234, _ := net.ResolveUDPAddr("udp4", "127.0.0.1:1234")
	global4Num1234, _ := net.ResolveUDPAddr("udp4", "1.2.3.4:1234")
	global4Num444, _ := net.ResolveUDPAddr("udp4", "1.2.3.4:444")
	local6Num666, _ := net.ResolveUDPAddr("udp6", "[::1]:666")
	local6Name666, err := net.ResolveUDPAddr("udp6", "localhost:666")
	if err != nil {
		local6Name666, _ = net.ResolveUDPAddr("udp6", "ip6-localhost:666")
	}
	local6Name1234, _ := net.ResolveUDPAddr("udp6", "localhost:1234")
	global6Num666, _ := net.ResolveUDPAddr("udp6", "[2a02::1234:1]:666")
	global6Num1234, _ := net.ResolveUDPAddr("udp6", "[2a02::1234:1]:1234")

	testvectors := []struct {
		a, b  *net.UDPAddr
		equal bool
	}{
		{nil, nil, true},
		{local4Num444, nil, false},
		{local4Num444, local4Num444, true},
		{local4Name444, local4Num444, true},
		{local4Name444, local4Num1234, false},
		{global4Num1234, local4Num1234, false},
		{global4Num1234, global4Num1234, true},
		{global4Num1234, global4Num444, false},
		{local6Num666, local6Num666, true},
		{local6Name666, local6Num666, true},
		{local6Name666, local6Name1234, false},
		{global6Num1234, local6Name1234, false},
		{global6Num1234, global6Num1234, true},
		{global6Num1234, global6Num666, false},
		{local4Num1234, local6Name1234, false},
	}

	for _, vector := range testvectors {
		if vector.equal {
			if !EndpointsEqual(vector.a, vector.b) {
				t.Fatalf("endpoints %v and %v should be equal but aren't", vector.a, vector.b)
			}
		} else {
			if EndpointsEqual(vector.a, vector.b) {
				t.Fatalf("endpoints %v and %v shouldn't be equal but are", vector.a, vector.b)
			}
		}
	}
}

func TestSecurityAssociationEndpointUpdate(t *testing.T) {
	addr4, _ := net.ResolveUDPAddr("udp4", "1.2.3.4:444")
	addr6, _ := net.ResolveUDPAddr("udp6", "[2a02::1]:666")

	// should this panic??
	sa := NewSecurityAssociation(nil, 0, 0, 0, 0)
	sa.EndpointUpdate(0, addr4)

	sa = NewSecurityAssociation(nil, 1, 0, 0, 0)
	if sa.endpoints[0] != nil {
		t.Fatalf("endpoints[0] is %v but should be nil", sa.endpoints[0])
	}
	sa.EndpointUpdate(0, addr4)
	if !EndpointsEqual(sa.endpoints[0], addr4) {
		t.Fatalf("endpoints[0] is %v but should be %v", sa.endpoints[0], addr4)
	}
	sa.EndpointUpdate(0, addr6)
	if !EndpointsEqual(sa.endpoints[0], addr6) {
		t.Fatalf("endpoints[0] is %v but should be %v", sa.endpoints[0], addr6)
	}

	sa = NewSecurityAssociation(nil, 3, 0, 0, 0)
	sa.EndpointUpdate(0, addr4)
	sa.EndpointUpdate(2, addr6)
	if !EndpointsEqual(sa.endpoints[0], addr4) {
		t.Fatalf("endpoints[0] is %v but should be %v", sa.endpoints[0], addr4)
	}
	if sa.endpoints[1] != nil {
		t.Fatalf("endpoints[1] is %v but should be nil", sa.endpoints[1])
	}
	if !EndpointsEqual(sa.endpoints[2], addr6) {
		t.Fatalf("endpoints[2] is %v but should be %v", sa.endpoints[2], addr6)
	}
}

func TestSecurityAssociationEndpointCompareAndUpdate(t *testing.T) {
	addr4, _ := net.ResolveUDPAddr("udp4", "1.2.3.4:444")
	addr6, _ := net.ResolveUDPAddr("udp6", "[2a01:1234::2]:666")

	// should this panic??
	sa := NewSecurityAssociation(nil, 0, 0, 0, 0)
	sa.EndpointCompareAndUpdate(0, addr4)

	sa = NewSecurityAssociation(nil, 1, 0, 0, 0)

	changed := sa.EndpointCompareAndUpdate(0, addr4)
	if !EndpointsEqual(sa.endpoints[0], addr4) {
		t.Fatalf("endpoints[0] is %v but should be %v", sa.endpoints[0], addr4)
	}
	if !changed {
		t.Fatal("updateting a nil endpoint should return changed = true")
	}

	changed = sa.EndpointCompareAndUpdate(0, addr4)
	if !EndpointsEqual(sa.endpoints[0], addr4) {
		t.Fatalf("endpoints[0] is %v but should be %v", sa.endpoints[0], addr4)
	}
	if changed {
		t.Fatal("updateting endpoint with same address should return changed = false")
	}

	changed = sa.EndpointCompareAndUpdate(0, addr6)
	if !EndpointsEqual(sa.endpoints[0], addr6) {
		t.Fatalf("endpoints[0] is %v but should be %v", sa.endpoints[0], addr6)
	}
	if !changed {
		t.Fatalf("updateting endpoint %v with %v should return changed = true", addr4, addr6)
	}
}

func TestSecurityAssociationGetEndpointsAndNextSequenceNumber(t *testing.T) {
	sa := NewSecurityAssociation(nil, 3, 0, 0, 0)

	seq, _ := sa.GetEndpointsAndNextSequenceNumber(nil)
	if seq != 0 {
		t.Fatalf("next sequnce number returned is %d but should be %d", seq, 0)
	}
	seq, _ = sa.GetEndpointsAndNextSequenceNumber(nil)
	seq, _ = sa.GetEndpointsAndNextSequenceNumber(nil)
	seq, _ = sa.GetEndpointsAndNextSequenceNumber(nil)
	if seq != 3 {
		t.Fatalf("next sequnce number returned is %d but should be %d", seq, 3)
	}

	sa = NewSecurityAssociation(nil, 3, (^uint32(0)), 0, 0)
	eps := make([]*net.UDPAddr, 3)
	for i := range eps {
		if eps[i] != nil {
			t.Fatalf("endpoints[%d] is %v but should be nil", i, eps[i])
		}
	}
	seq, _ = sa.GetEndpointsAndNextSequenceNumber(eps)
	if seq != (^uint32(0)) {
		t.Fatalf("next sequnce number returned is %d but should be %d", seq, (^uint32(0)))
	}
	addr4, _ := net.ResolveUDPAddr("udp4", "1.2.3.4:444")
	sa.EndpointUpdate(0, addr4)
	addr6, _ := net.ResolveUDPAddr("udp6", "[2a01:1234::2]:666")
	sa.EndpointUpdate(2, addr6)
	seq, _ = sa.GetEndpointsAndNextSequenceNumber(eps)
	if seq != 0 {
		t.Fatalf("next sequnce number returned is %d but should be %d", seq, 0)
	}
	if !EndpointsEqual(sa.endpoints[0], addr4) {
		t.Fatalf("endpoints[0] is %v but should be %v", sa.endpoints[0], addr4)
	}
	if sa.endpoints[1] != nil {
		t.Fatalf("endpoints[1] is %v but should be nil", sa.endpoints[1])
	}
	if !EndpointsEqual(sa.endpoints[2], addr6) {
		t.Fatalf("endpoints[2] is %v but should be %v", sa.endpoints[2], addr6)
	}
}

func TestSecurityAssociationSequenceNumberCheck(t *testing.T) {
	sa := NewSecurityAssociation(nil, 1, 0, 23, 10)
	if sa.SequenceNumberCheck(0, 0) {
		t.Fatal("sequence number 0 from sender 0 shouldn't get accepted")
	}
	if !sa.SequenceNumberCheck(0, 23) {
		t.Fatal("sequence number 23 from sender 0 should get accepted")
	}
	if !sa.SequenceNumberCheckAndSet(0, 23) {
		t.Fatal("sequence number 23 from sender 0 should get accepted")
	}
	if sa.SequenceNumberCheck(0, 23) {
		t.Fatal("sequence number 23 from sender 0 shouldn't get accepted")
	}
	if !sa.SequenceNumberCheck(42, 23) {
		t.Fatal("sequence number 23 from sender 42 should get accepted")
	}
	if !sa.SequenceNumberCheckAndSet(42, 23) {
		t.Fatal("sequence number 23 from sender 42 should get accepted")
	}
	if sa.SequenceNumberCheck(42, 23) {
		t.Fatal("sequence number 23 from sender 42 shouldn't get accepted")
	}
	if !sa.SequenceNumberCheckAndSet(23, 27) {
		t.Fatal("sequence number 27 from sender 23 should get accepted")
	}

	sa.seqWindows.Range(func(key, value interface{}) bool {
		t.Logf("SeqWin for Sender %v: %v", key, value)
		return true
	})
}