package aerotun

import (
	"bytes"
	"encoding/binary"
	"errors"
	"net/netip"
	"sync"
	"testing"
	"time"
)

func secureConfigs(t testing.TB) (TunnelInstance, TunnelInstance) {
	t.Helper()
	privA, pubA, err := GenerateIdentity()
	if err != nil {
		t.Fatal(err)
	}
	privB, pubB, err := GenerateIdentity()
	if err != nil {
		t.Fatal(err)
	}
	a := TunnelInstance{Protocol: 3, Mode: "client", PSK: testPSK, PrivateKey: privA, PublicKey: pubA, PeerPublicKey: pubB, Name: "atest0", LocalAddress: "10.99.0.2/32", PeerAddress: "10.99.0.1/32", AllowedIPs: []string{"10.99.0.1/32"}, RemoteEndpoint: "192.0.2.1:51830", MTU: 1420}
	b := a
	b.Mode = "server"
	b.PrivateKey = privB
	b.PublicKey = pubB
	b.PeerPublicKey = pubA
	b.LocalAddress, b.PeerAddress = a.PeerAddress, a.LocalAddress
	b.AllowedIPs = []string{"10.99.0.2/32"}
	b.RemoteEndpoint = ""
	return a, b
}
func securePair(t testing.TB) (*secureProtocol, *secureProtocol) {
	t.Helper()
	a, b := secureConfigs(t)
	pa, err := newSecureProtocol(a)
	if err != nil {
		t.Fatal(err)
	}
	pb, err := newSecureProtocol(b)
	if err != nil {
		t.Fatal(err)
	}
	return pa, pb
}
func secureEstablish(t testing.TB, a, b *secureProtocol) {
	t.Helper()
	b.mu.Lock()
	b.lastRequest = time.Time{}
	b.mu.Unlock()
	hello, err := a.initiation(1420, true)
	if err != nil || len(hello) == 0 {
		t.Fatalf("initiation: %v", err)
	}
	challenge, _, err := b.handleHandshake(hello, 1420)
	if err != nil || challenge[3] != secureCookie {
		t.Fatal("cookie challenge failed", err)
	}
	request, _, err := a.handleHandshake(challenge, 1420)
	if err != nil {
		t.Fatal(err)
	}
	response, _, err := b.handleHandshake(request, 1420)
	if err != nil {
		t.Fatal(err)
	}
	if _, mtu, err := a.handleHandshake(response, 1420); err != nil || mtu != 1420 {
		t.Fatalf("response %v %d", err, mtu)
	}
	var buf [128]byte
	var plain [128]byte
	var txNonce, rxNonce [24]byte
	packet, err := a.sealInto(buf[:0], &txNonce, FlagKeepalive, []byte{5, 140})
	if err != nil {
		t.Fatal(err)
	}
	if _, _, mtu, err := b.openInto(plain[:0], &rxNonce, packet); err != nil || (mtu != 1420 && mtu != 0) {
		t.Fatalf("confirmation: %v %d", err, mtu)
	}
}
func TestSecureHandshakeReplayPinsAndRekey(t *testing.T) {
	a, b := securePair(t)
	secureEstablish(t, a, b)
	var txBuf, rxBuf [2048]byte
	var tn, rn [24]byte
	raw := []byte("authenticated application payload")
	packet, err := a.sealInto(txBuf[:0], &tn, 0, raw)
	if err != nil {
		t.Fatal(err)
	}
	oldPacket := append([]byte(nil), packet...)
	forged := append([]byte(nil), packet...)
	binary.BigEndian.PutUint64(forged[8:16], 1<<40)
	if _, _, _, err = b.openInto(rxBuf[:0], &rn, forged); err == nil {
		t.Fatal("forged high counter accepted")
	}
	got, _, _, err := b.openInto(rxBuf[:0], &rn, packet)
	if err != nil || !bytes.Equal(got, raw) {
		t.Fatal("valid packet lost after forgery", err)
	}
	if _, _, _, err = b.openInto(rxBuf[:0], &rn, packet); !errors.Is(err, errReplay) {
		t.Fatalf("replay accepted: %v", err)
	}
	if _, _, _, err = a.openInto(rxBuf[:0], &rn, packet); err == nil {
		t.Fatal("reflected packet accepted")
	}
	oldSession := a.sessions.Load().current
	secureEstablish(t, a, b)
	if a.sessions.Load().current == oldSession || a.sessions.Load().current.txPrefix == oldSession.txPrefix {
		t.Fatal("rekey reused data key/nonce context")
	}
	if _, _, _, err = b.openInto(rxBuf[:0], &rn, oldPacket); !errors.Is(err, errReplay) {
		t.Fatal("old duplicate accepted during grace", err)
	}
	oldB := b.sessions.Load()
	b.sessions.Store(&secureSessions{current: oldB.current, previous: oldB.previous, previousUntil: time.Now().Add(-time.Second)})
	if _, _, _, err = b.openInto(rxBuf[:0], &rn, oldPacket); err == nil {
		t.Fatal("expired session accepted")
	}
	// An independent boot has fresh ephemeral keys and rejects old recordings.
	restarted, err := newSecureProtocol(b.cfg)
	if err != nil {
		t.Fatal(err)
	}
	if _, _, _, err = restarted.openInto(rxBuf[:0], &rn, oldPacket); err == nil {
		t.Fatal("old packet survived receiver restart")
	}
	secureEstablish(t, a, restarted)
}
func TestSecureWrongPeerAndWrongPSK(t *testing.T) {
	a, b := securePair(t)
	_, wrongPub, err := GenerateIdentity()
	if err != nil {
		t.Fatal(err)
	}
	b.peer, _ = decodeKey(wrongPub)
	packet, err := a.initiation(1420, true)
	if err != nil {
		t.Fatal(err)
	}
	challenge, _, err := b.handleHandshake(packet, 1420)
	if err != nil {
		t.Fatal(err)
	}
	request, _, err := a.handleHandshake(challenge, 1420)
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err = b.handleHandshake(request, 1420); err == nil {
		t.Fatal("unpinned peer accepted with correct PSK")
	}
	a, b = securePair(t)
	b.psk = bytes.Repeat([]byte{42}, 32)
	packet, _ = a.initiation(1420, true)
	if _, _, err = b.handleHandshake(packet, 1420); err == nil {
		t.Fatal("wrong PSK accepted")
	}
}
func TestSecureDuplicateInitiationDoesNotActivateSession(t *testing.T) {
	a, b := securePair(t)
	init, _ := a.initiation(1420, true)
	challenge, _, err := b.handleHandshake(init, 1420)
	if err != nil {
		t.Fatal(err)
	}
	request, _, err := a.handleHandshake(challenge, 1420)
	if err != nil {
		t.Fatal(err)
	}
	response, _, err := b.handleHandshake(request, 1420)
	if err != nil {
		t.Fatal(err)
	}
	if b.sessions.Load().current != nil {
		t.Fatal("responder activated before key confirmation")
	}
	if _, _, err = b.handleHandshake(request, 1420); !errors.Is(err, errHandshakeRate) {
		t.Fatal("no handshake rate limit")
	}
	b.lastRequest = time.Now().Add(-2 * time.Second)
	again, _, err := b.handleHandshake(request, 1420)
	if err != nil || !bytes.Equal(response, again) {
		t.Fatal("retry regenerated keys", err)
	}
}

func TestSecureConcurrentCounters(t *testing.T) {
	a, b := securePair(t)
	secureEstablish(t, a, b)
	var wg sync.WaitGroup
	counts := make(chan uint64, 400)
	for range 4 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			var buf [128]byte
			var nonce [24]byte
			for range 100 {
				out, err := a.sealInto(buf[:0], &nonce, 0, []byte("data"))
				if err != nil {
					t.Error(err)
					return
				}
				counts <- binary.BigEndian.Uint64(out[8:16])
			}
		}()
	}
	wg.Wait()
	close(counts)
	seen := map[uint64]bool{}
	for n := range counts {
		if seen[n] {
			t.Fatal("nonce counter reused")
		}
		seen[n] = true
	}
}
func TestSecureReplayReordering(t *testing.T) {
	var r secureReplay
	for _, n := range []uint64{1, 1024, 2, 1023, 2048, 2047} {
		if !r.accept(n) {
			t.Fatal("valid reordered packet rejected", n)
		}
	}
	for _, n := range []uint64{0, 1, 2, 1024, 2047, 2048} {
		if r.accept(n) {
			t.Fatal("old/duplicate accepted", n)
		}
	}
}
func BenchmarkSecureProtocolRoundTrip(b *testing.B) {
	a, peer := securePair(b)
	secureEstablish(b, a, peer)
	tx := make([]byte, 1500)
	rx := make([]byte, 1500)
	payload := make([]byte, 1200)
	var tn, rn [24]byte
	b.SetBytes(1200)
	b.ReportAllocs()
	b.ResetTimer()
	for range b.N {
		packet, err := a.sealInto(tx[:0], &tn, 0, payload)
		if err != nil {
			b.Fatal(err)
		}
		if _, _, _, err = peer.openInto(rx[:0], &rn, packet); err != nil {
			b.Fatal(err)
		}
	}
}
func FuzzSecurePacket(f *testing.F) {
	a, b := securePair(f)
	secureEstablish(f, a, b)
	f.Add([]byte("invalid"))
	var wire [128]byte
	var nonce [24]byte
	valid, _ := a.sealInto(wire[:0], &nonce, 0, []byte("payload"))
	f.Add(append([]byte(nil), valid...))
	f.Fuzz(func(t *testing.T, packet []byte) {
		var dst [2048]byte
		var n [24]byte
		_, _, _, _ = b.openInto(dst[:0], &n, packet)
		_, _, _ = b.handleHandshake(packet, 1420)
	})
}

func TestSecureInPlaceSealAndSourcePolicy(t *testing.T) {
	a, b := securePair(t)
	secureEstablish(t, a, b)
	wire := make([]byte, 1500)
	payload := wire[16:1216]
	for i := range payload {
		payload[i] = byte(i)
	}
	expected := append([]byte(nil), payload...)
	var tn, rn [24]byte
	var plain [1500]byte
	packet, err := a.sealInto(wire[:0], &tn, 0, payload)
	if err != nil {
		t.Fatal(err)
	}
	got, _, _, err := b.openInto(plain[:0], &rn, packet)
	if err != nil || !bytes.Equal(got, expected) {
		t.Fatal("in-place encryption corrupted packet", err)
	}
}
func TestSecureRekeyUnderTraffic(t *testing.T) {
	a, b := securePair(t)
	secureEstablish(t, a, b)
	var wg sync.WaitGroup
	started := make(chan struct{})
	wg.Add(1)
	go func() {
		defer wg.Done()
		var wire, plain [512]byte
		var tn, rn [24]byte
		close(started)
		for range 1000 {
			out, err := a.sealInto(wire[:0], &tn, 0, []byte("traffic during key rotation"))
			if err != nil {
				t.Error(err)
				return
			}
			if _, _, _, err = b.openInto(plain[:0], &rn, out); err != nil {
				t.Error(err)
				return
			}
		}
	}()
	<-started
	secureEstablish(t, a, b)
	wg.Wait()
	if a.handshakes.Load() != 2 || b.handshakes.Load() != 2 {
		t.Fatal("rotation not completed")
	}
}

func TestSecureCookieBindsSourceAndExpires(t *testing.T) {
	a, b := securePair(t)
	init, _ := a.initiation(1420, true)
	address := netip.MustParseAddrPort("192.0.2.10:51830")
	other := netip.MustParseAddrPort("192.0.2.11:51830")
	challenge, _, err := b.handleHandshake(init, 1420, address)
	if err != nil || challenge[3] != secureCookie {
		t.Fatal("missing source cookie", err)
	}
	request, _, err := a.handleHandshake(challenge, 1420)
	if err != nil {
		t.Fatal(err)
	}
	reply, _, err := b.handleHandshake(request, 1420, other)
	if err != nil || reply[3] != secureCookie || b.sessions.Load().pending != nil {
		t.Fatal("cookie accepted from another source", err)
	}
	expired := append([]byte(nil), request...)
	old := b.addressCookie(expired[:16], address, time.Now().Unix()/30-3)
	copy(expired[len(expired)-32:len(expired)-16], old[:])
	expired = a.finishHandshake(expired[:len(expired)-16])
	reply, _, err = b.handleHandshake(expired, 1420, address)
	if err != nil || reply[3] != secureCookie || b.sessions.Load().pending != nil {
		t.Fatal("expired cookie accepted", err)
	}
	if reply, _, err = b.handleHandshake(request, 1420, address); err != nil || reply[3] != secureResponse {
		t.Fatal("valid source cookie rejected", err)
	}
}

func TestSecureAllowedSourcePolicy(t *testing.T) {
	runner := TunnelRunner{allowed: []netip.Prefix{netip.MustParsePrefix("172.16.233.2/32"), netip.MustParsePrefix("2001:db8:42::/64")}}
	packet := make([]byte, 20)
	packet[0] = 0x45
	packet[3] = 20
	for _, tc := range []struct {
		source string
		want   bool
	}{{"172.16.233.2", true}, {"172.16.233.1", false}, {"192.168.1.1", false}} {
		source := netip.MustParseAddr(tc.source).As4()
		copy(packet[12:16], source[:])
		if runner.allowedSource(packet) != tc.want {
			t.Fatalf("incorrect IPv4 source policy for %s", tc.source)
		}
	}
	v6 := make([]byte, 40)
	v6[0] = 0x60
	for _, tc := range []struct {
		source string
		want   bool
	}{{"2001:db8:42::1", true}, {"2001:db8:43::1", false}, {"fe80::1423:58a:b322:ae3a", true}} {
		source := netip.MustParseAddr(tc.source).As16()
		copy(v6[8:24], source[:])
		if runner.allowedSource(v6) != tc.want {
			t.Fatalf("incorrect IPv6 source policy for %s", tc.source)
		}
	}
}