package peer

import (
	"bytes"
	"reflect"
	"testing"
)

func TestRemotePeer_DecryptDataPacket(t *testing.T) {
	p1, p2, _ := NewPeersForTesting()
	orig := RandPacket()

	peer2 := p1.RT.Load().Peers[2]
	peer1 := p2.RT.Load().Peers[1]

	enc := peer2.EncryptDataPacket(2, orig, newBuf())

	h := parseHeader(enc)
	if h.DestIP != 2 || h.SourceIP != 1 {
		t.Fatal(h)
	}

	dec, err := peer1.DecryptDataPacket(h, enc, newBuf())
	if err != nil {
		t.Fatal(err)
	}

	if !bytes.Equal(orig, dec) {
		t.Fatal(dec)
	}
}

func TestRemotePeer_DecryptDataPacket_packetAltered(t *testing.T) {
	p1, p2, _ := NewPeersForTesting()
	orig := RandPacket()

	peer2 := p1.RT.Load().Peers[2]
	peer1 := p2.RT.Load().Peers[1]

	enc := peer2.EncryptDataPacket(2, orig, newBuf())

	h := parseHeader(enc)

	for range 2048 {
		_, err := peer1.DecryptDataPacket(h, ModifyPacket(enc), newBuf())
		if err == nil {
			t.Fatal(enc)
		}
	}
}

func TestRemotePeer_DecryptDataPacket_duplicateSequenceNumber(t *testing.T) {
	p1, p2, _ := NewPeersForTesting()
	orig := RandPacket()

	peer2 := p1.RT.Load().Peers[2]
	peer1 := p2.RT.Load().Peers[1]

	enc := peer2.EncryptDataPacket(2, orig, newBuf())
	h := parseHeader(enc)

	if _, err := peer1.DecryptDataPacket(h, enc, newBuf()); err != nil {
		t.Fatal(err)
	}

	if _, err := peer1.DecryptDataPacket(h, enc, newBuf()); err == nil {
		t.Fatal(err)
	}
}

func TestRemotePeer_DecryptControlPacket(t *testing.T) {
	p1, p2, _ := NewPeersForTesting()

	peer2 := p1.RT.Load().Peers[2]
	peer1 := p2.RT.Load().Peers[1]

	orig := PacketProbe{TraceID: newTraceID()}

	enc := peer2.EncryptControlPacket(orig, newBuf(), newBuf())

	h := parseHeader(enc)
	if h.DestIP != 2 || h.SourceIP != 1 {
		t.Fatal(h)
	}

	ctrlMsg, err := peer1.DecryptControlPacket(p1.RT.Load().LocalAddr, h, enc, newBuf())
	if err != nil {
		t.Fatal(err)
	}

	dec, ok := ctrlMsg.(controlMsg[PacketProbe])
	if !ok {
		t.Fatal(ctrlMsg)
	}

	if dec.SrcIP != 1 || dec.SrcAddr != p1.RT.Load().LocalAddr {
		t.Fatal(dec)
	}

	if !reflect.DeepEqual(dec.Packet, orig) {
		t.Fatal(dec)
	}
}

func TestRemotePeer_DecryptControlPacket_packetAltered(t *testing.T) {
	p1, p2, _ := NewPeersForTesting()

	peer2 := p1.RT.Load().Peers[2]
	peer1 := p2.RT.Load().Peers[1]

	orig := PacketProbe{TraceID: newTraceID()}

	enc := peer2.EncryptControlPacket(orig, newBuf(), newBuf())

	h := parseHeader(enc)
	if h.DestIP != 2 || h.SourceIP != 1 {
		t.Fatal(h)
	}

	for range 2048 {
		ctrlMsg, err := peer1.DecryptControlPacket(p1.RT.Load().LocalAddr, h, ModifyPacket(enc), newBuf())
		if err == nil {
			t.Fatal(ctrlMsg)
		}
	}
}

func TestRemotePeer_DecryptControlPacket_duplicateSequenceNumber(t *testing.T) {
	p1, p2, _ := NewPeersForTesting()

	peer2 := p1.RT.Load().Peers[2]
	peer1 := p2.RT.Load().Peers[1]

	orig := PacketProbe{TraceID: newTraceID()}

	enc := peer2.EncryptControlPacket(orig, newBuf(), newBuf())

	h := parseHeader(enc)
	if h.DestIP != 2 || h.SourceIP != 1 {
		t.Fatal(h)
	}

	if _, err := peer1.DecryptControlPacket(p1.RT.Load().LocalAddr, h, enc, newBuf()); err != nil {
		t.Fatal(err)
	}
	if _, err := peer1.DecryptControlPacket(p1.RT.Load().LocalAddr, h, enc, newBuf()); err == nil {
		t.Fatal(err)
	}
}

func TestRemotePeer_DecryptControlPacket_unknownPacketType(t *testing.T) {
	p1, p2, _ := NewPeersForTesting()

	peer2 := p1.RT.Load().Peers[2]
	peer1 := p2.RT.Load().Peers[1]

	orig := UnknownControlPacket{TraceID: newTraceID()}

	enc := peer2.EncryptControlPacket(orig, newBuf(), newBuf())

	h := parseHeader(enc)
	if h.DestIP != 2 || h.SourceIP != 1 {
		t.Fatal(h)
	}

	if _, err := peer1.DecryptControlPacket(p1.RT.Load().LocalAddr, h, enc, newBuf()); err == nil {
		t.Fatal(err)
	}
}