package nebula

import (
	"math/rand/v2"
	"net/netip"
	"sync"
	"testing"
	"time"

	"github.com/slackhq/nebula/config"
	"github.com/slackhq/nebula/handshake/fragment"
	"github.com/slackhq/nebula/header"
	"github.com/slackhq/nebula/test"
	"github.com/slackhq/nebula/udp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// recordingConn captures what is written to it.
type recordingConn struct {
	udp.NoopConn
	sync.Mutex
	sent [][]byte
}

func (c *recordingConn) WriteTo(b []byte, _ netip.AddrPort) error {
	c.Lock()
	defer c.Unlock()
	c.sent = append(c.sent, append([]byte(nil), b...))
	return nil
}

func (c *recordingConn) take() [][]byte {
	c.Lock()
	defer c.Unlock()
	s := c.sent
	c.sent = nil
	return s
}

func newPQTestManager(t *testing.T) (*HandshakeManager, *recordingConn) {
	l := test.NewLogger()
	conn := &recordingConn{}
	cfg := defaultHandshakeConfig
	cfg.pq = pqHandshakeConfig{initiate: true, accept: true}
	hm := NewHandshakeManager(l, newHostMap(l), newTestLighthouse(), conn, cfg)
	hm.f = &Interface{
		handshakeManager: hm,
		outside:          conn,
		messageMetrics:   newMessageMetricsOnlyRecvError(),
		l:                l,
	}
	return hm, conn
}

// fragmentIndexes checks every packet is a fragment of the expected response and
// returns their symbol indexes.
func fragmentIndexes(t *testing.T, pkts [][]byte, fragID, remoteIndex uint32) []int {
	t.Helper()
	var out []int
	for _, p := range pkts {
		require.LessOrEqual(t, len(p), fragment.MaxPacketLen)
		var h header.H
		require.NoError(t, h.Parse(p))
		assert.Equal(t, header.HandshakeIXPSK2MLKEM1024, h.Subtype)
		assert.Equal(t, remoteIndex, h.RemoteIndex)
		assert.Equal(t, uint64(2), h.MessageCounter)
		fh, _, err := fragment.Parse(p[header.Len:])
		require.NoError(t, err)
		assert.Equal(t, fragID, fh.FragID)
		out = append(out, int(fh.Index))
	}
	return out
}

func TestSendFragmentedResponse(t *testing.T) {
	hm, conn := newPQTestManager(t)
	via := ViaSender{UdpAddr: netip.MustParseAddrPort("192.0.2.1:4242")}
	hostinfo := &HostInfo{localIndexId: 77}

	// A 2000 byte message is 2 symbols of 1000, so packets of 1024 bytes.
	msg := make([]byte, header.Len, header.Len+2000)
	header.Encode(msg, header.Version, header.Handshake, header.HandshakeIXPSK2MLKEM1024, 5, 2)
	r := rand.New(rand.NewPCG(1, 1))
	for range 2000 {
		msg = append(msg, byte(r.Uint32()))
	}
	const pktLen = header.Len + fragment.HeaderLen + 1000

	t.Run("first send", func(t *testing.T) {
		require.NoError(t, hm.sendFragmentedResponse(via, msg, hostinfo, false, 0))
		assert.Equal(t, []int{0, 1, 2, 3}, fragmentIndexes(t, conn.take(), 77, 5))
	})

	t.Run("resends are rate limited", func(t *testing.T) {
		require.Error(t, hm.sendFragmentedResponse(via, msg, hostinfo, true, 10*pktLen))
		assert.Empty(t, conn.take())
	})

	allowResend := func() {
		hostinfo.handshakeResendAt.Store(time.Now().Add(-time.Second).UnixNano())
	}

	t.Run("resend offers fresh symbols", func(t *testing.T) {
		allowResend()
		require.NoError(t, hm.sendFragmentedResponse(via, msg, hostinfo, true, 10*pktLen))
		assert.Equal(t, []int{4, 5, 6, 7}, fragmentIndexes(t, conn.take(), 77, 5))
	})

	t.Run("amplification budget trims repair", func(t *testing.T) {
		allowResend()
		// 3x of a little over one packet allows 3 packets: both data and 1 repair.
		require.NoError(t, hm.sendFragmentedResponse(via, msg, hostinfo, true, pktLen+10))
		assert.Equal(t, []int{8, 9, 10}, fragmentIndexes(t, conn.take(), 77, 5))
	})

	t.Run("amplification budget below the data refuses", func(t *testing.T) {
		allowResend()
		require.Error(t, hm.sendFragmentedResponse(via, msg, hostinfo, true, pktLen/2))
		assert.Empty(t, conn.take())
	})
}

func TestPQAnswered(t *testing.T) {
	now := time.Unix(1000, 0)
	a := newPQAnswered(5 * time.Second)
	hi := &HostInfo{}

	assert.Nil(t, a.lookup([]byte("msg1"), now))
	a.add([]byte("msg1"), hi, now)
	assert.Same(t, hi, a.lookup([]byte("msg1"), now.Add(time.Second)))
	assert.Nil(t, a.lookup([]byte("other"), now))
	assert.Nil(t, a.lookup([]byte("msg1"), now.Add(6*time.Second)), "stale entries are not answered")

	a.expire(now.Add(6 * time.Second))
	assert.Empty(t, a.entries)

	for i := range pqAnsweredMax + 10 {
		a.add([]byte{byte(i), byte(i >> 8)}, hi, now)
	}
	assert.Len(t, a.entries, pqAnsweredMax)
}

func TestPQFragKeyFor(t *testing.T) {
	k1, s1 := pqFragKeyFor(ViaSender{UdpAddr: netip.MustParseAddrPort("192.0.2.1:4242")}, 9)
	k2, s2 := pqFragKeyFor(ViaSender{UdpAddr: netip.MustParseAddrPort("192.0.2.1:4243")}, 9)
	assert.NotEqual(t, k1, k2, "different ports are different messages")
	assert.Equal(t, s1, s2, "but share an IPv4 /32 quota")

	_, s3 := pqFragKeyFor(ViaSender{UdpAddr: netip.MustParseAddrPort("[2001:db8::1]:1")}, 9)
	_, s4 := pqFragKeyFor(ViaSender{UdpAddr: netip.MustParseAddrPort("[2001:db8::ffff]:1")}, 9)
	assert.Equal(t, s3, s4, "an IPv6 /64 shares a quota")

	_, s5 := pqFragKeyFor(ViaSender{UdpAddr: netip.MustParseAddrPort("[::ffff:192.0.2.1]:1")}, 9)
	assert.Equal(t, s1, s5, "v4-mapped addresses count as IPv4")
}

func TestPQCookies(t *testing.T) {
	now := time.Unix(1000, 0)
	c := newPQCookies(now)
	key := pqFragKey{addr: netip.MustParseAddrPort("192.0.2.1:4242"), fragID: 9}

	cookie := c.make(key, now)
	assert.True(t, c.valid(key, cookie, now))

	for _, other := range []pqFragKey{
		{addr: netip.MustParseAddrPort("192.0.2.2:4242"), fragID: 9},
		{addr: netip.MustParseAddrPort("192.0.2.1:4243"), fragID: 9},
		{addr: netip.MustParseAddrPort("192.0.2.1:4242"), fragID: 10},
		{addr: netip.MustParseAddrPort("192.0.2.1:4242"), relay: 1, fragID: 9},
	} {
		assert.False(t, c.valid(other, cookie, now), "%+v", other)
	}

	now = now.Add(pqCookieRotate)
	assert.True(t, c.valid(key, cookie, now), "the previous secret is still accepted")
	assert.NotEqual(t, cookie, c.make(key, now), "new cookies come from the new secret")

	now = now.Add(pqCookieRotate)
	assert.False(t, c.valid(key, cookie, now), "two rotations retire a cookie")
}

// pqMsg1Fragment builds the first fragment of a msg1, optionally with a cookie.
func pqMsg1Fragment(t *testing.T, fragID uint32, cookie *fragment.Cookie) ([]byte, *header.H) {
	t.Helper()
	s, err := fragment.NewSender(make([]byte, 2000), header.HandshakeIXPSK2MLKEM1024, 0, 1, fragID)
	require.NoError(t, err)
	if cookie != nil {
		s.SetCookie(*cookie)
	}
	pkts, _, err := s.Attempt(0)
	require.NoError(t, err)
	var h header.H
	require.NoError(t, h.Parse(pkts[0]))
	return pkts[0], &h
}

// fillPQTable puts the responder's msg1 table at the cookie threshold.
func fillPQTable(t *testing.T, hm *HandshakeManager) {
	t.Helper()
	now := time.Now()
	for i := 0; hm.pqTable.Load() < pqCookieLoad; i++ {
		p, _ := pqMsg1Fragment(t, uint32(i+1), nil)
		fh, sym, err := fragment.Parse(p[header.Len:])
		require.NoError(t, err)
		addr := netip.AddrPortFrom(netip.AddrFrom4([4]byte{198, 18, byte(i >> 8), byte(i)}), 4242)
		key, src := pqFragKeyFor(ViaSender{UdpAddr: addr}, fh.FragID)
		_, err = hm.pqTable.Add(key, src, now, fh, sym)
		require.NoError(t, err)
	}
}

func TestPQCookieUnderLoad(t *testing.T) {
	hm, conn := newPQTestManager(t)
	via := ViaSender{UdpAddr: netip.MustParseAddrPort("192.0.2.1:4242")}

	t.Run("no cookie needed below the threshold", func(t *testing.T) {
		p, h := pqMsg1Fragment(t, 5000, nil)
		before := hm.pqTable.Len()
		hm.handleFragment(via, p, h)
		assert.Empty(t, conn.take())
		assert.Equal(t, before+1, hm.pqTable.Len())
	})

	fillPQTable(t, hm)

	var cookie fragment.Cookie
	t.Run("a fragment without a cookie gets one and leaves no state", func(t *testing.T) {
		p, h := pqMsg1Fragment(t, 7, nil)
		before := hm.pqTable.Len()
		hm.handleFragment(via, p, h)
		assert.Equal(t, before, hm.pqTable.Len())

		sent := conn.take()
		require.Len(t, sent, 1)
		reply := sent[0]
		require.Len(t, reply, pqCookieReplyLen)
		var rh header.H
		require.NoError(t, rh.Parse(reply))
		assert.Equal(t, header.Handshake, rh.Type)
		assert.Equal(t, header.HandshakeIXPSK2MLKEM1024, rh.Subtype)
		assert.Equal(t, uint32(7), rh.RemoteIndex, "the reply names the initiator's FragID")
		assert.Equal(t, uint64(pqCookieCounter), rh.MessageCounter)
		copy(cookie[:], reply[header.Len:])
	})

	t.Run("a forged cookie gets a fresh one", func(t *testing.T) {
		bad := cookie
		bad[0] ^= 1
		p, h := pqMsg1Fragment(t, 7, &bad)
		before := hm.pqTable.Len()
		hm.handleFragment(via, p, h)
		assert.Equal(t, before, hm.pqTable.Len())
		assert.Len(t, conn.take(), 1)
	})

	t.Run("the cookie gets the fragment in", func(t *testing.T) {
		p, h := pqMsg1Fragment(t, 7, &cookie)
		before := hm.pqTable.Len()
		hm.handleFragment(via, p, h)
		assert.Equal(t, before+1, hm.pqTable.Len())
		assert.Empty(t, conn.take())
	})

	t.Run("a cookie is bound to its address", func(t *testing.T) {
		p, h := pqMsg1Fragment(t, 7, &cookie)
		hm.handleFragment(ViaSender{UdpAddr: netip.MustParseAddrPort("192.0.2.99:4242")}, p, h)
		assert.Len(t, conn.take(), 1, "another address gets its own cookie instead")
	})

	t.Run("a reply never outgrows what prompted it", func(t *testing.T) {
		p, h := pqMsg1Fragment(t, 8, nil)
		hm.sendCookie(via, header.HandshakeIXPSK2MLKEM1024, fragment.Header{FragID: 8}, pqFragKey{}, pqCookieReplyLen-1, time.Now())
		assert.Empty(t, conn.take())
		hm.handleFragment(via, p, h)
		assert.Len(t, conn.take(), 1)
	})
}

func TestPQHandleCookie(t *testing.T) {
	hm, conn := newPQTestManager(t)
	via := ViaSender{UdpAddr: netip.MustParseAddrPort("192.0.2.1:4242")}

	s, err := fragment.NewSender(make([]byte, 2000), header.HandshakeIXPSK2MLKEM1024, 0, 1, 99)
	require.NoError(t, err)
	_, _, err = s.Attempt(2)
	require.NoError(t, err)
	hh := &HandshakeHostInfo{hostinfo: &HostInfo{}, pq: &pqAttempt{index: 99, sender: s}}
	hm.indexes[99] = hh

	reply := make([]byte, header.Len, pqCookieReplyLen)
	header.Encode(reply, header.Version, header.Handshake, header.HandshakeIXPSK2MLKEM1024, 99, pqCookieCounter)
	reply = append(reply, bytes16(0xab)...)
	var h header.H
	require.NoError(t, h.Parse(reply))

	hm.handleFragment(via, reply, &h)
	assert.True(t, hh.pq.heard, "a cookie reply shows the peer speaks PQ")
	sent := conn.take()
	require.Len(t, sent, 4, "an immediate resend of fresh symbols")
	for i, p := range sent {
		fh, _, err := fragment.Parse(p[header.Len:])
		require.NoError(t, err)
		assert.True(t, fh.HasCookie())
		assert.Equal(t, fragment.Cookie(bytes16(0xab)), fh.Cookie)
		assert.Equal(t, uint8(4+i), fh.Index)
	}

	hm.handleFragment(via, reply, &h)
	assert.Empty(t, conn.take(), "immediate resends are rate limited")

	t.Run("replies for other indexes are ignored", func(t *testing.T) {
		other := append([]byte(nil), reply...)
		header.Encode(other, header.Version, header.Handshake, header.HandshakeIXPSK2MLKEM1024, 98, pqCookieCounter)
		var oh header.H
		require.NoError(t, oh.Parse(other))
		hh.pq.cookieResendAt = time.Time{}
		hm.handleFragment(via, other, &oh)
		assert.Empty(t, conn.take())
	})
}

func bytes16(b byte) []byte {
	out := make([]byte, 16)
	for i := range out {
		out[i] = b
	}
	return out
}

func TestReloadPQ(t *testing.T) {
	l := test.NewLogger()
	c := config.NewC(l)
	require.NoError(t, c.LoadString(`handshakes: {pq: {initiate: false}}`))
	cfg := defaultHandshakeConfig
	cfg.pq = pqHandshakeConfigFrom(c)
	hm := NewHandshakeManager(l, newHostMap(l), newTestLighthouse(), &udp.NoopConn{}, cfg)
	c.RegisterReloadCallback(hm.reloadPQ)

	assert.Equal(t, pqHandshakeConfig{repairDelay: DefaultPQRepairDelay}, *hm.pqConfig())

	require.NoError(t, c.ReloadConfigString(`handshakes: {pq: {initiate: true, accept: true, sticky: true, repair_delay: 25ms}}`))
	assert.Equal(t, pqHandshakeConfig{initiate: true, accept: true, sticky: true, repairDelay: 25 * time.Millisecond}, *hm.pqConfig())

	before := hm.pqConfig()
	require.NoError(t, c.ReloadConfigString(`handshakes: {try_interval: 1s, pq: {initiate: true, accept: true, sticky: true, repair_delay: 25ms}}`))
	assert.Same(t, before, hm.pqConfig(), "an unrelated change leaves the settings alone")

	require.NoError(t, c.ReloadConfigString(`handshakes: {}`))
	assert.Equal(t, pqHandshakeConfig{repairDelay: DefaultPQRepairDelay}, *hm.pqConfig(), "removing the section restores the defaults")

	require.NoError(t, c.ReloadConfigString(`handshakes: {pq: {require: true}}`))
	assert.Equal(t, pqHandshakeConfig{initiate: true, accept: true, require: true, repairDelay: DefaultPQRepairDelay}, *hm.pqConfig(), "require implies initiate and accept")
}

func TestUsePQ(t *testing.T) {
	hm, _ := newPQTestManager(t)
	peer := netip.MustParseAddr("10.1.0.1")
	set := func(cfg pqHandshakeConfig) { hm.pq.Store(&cfg) }

	// attempts plays a run of attempts, given whether each was triggered early,
	// and returns which were PQ.
	attempts := func(triggered ...bool) []bool {
		hh := &HandshakeHostInfo{}
		var out []bool
		for _, tr := range triggered {
			hh.counter++
			hh.lastPQ = hm.usePQ(hh, peer, tr)
			out = append(out, hh.lastPQ)
		}
		return out
	}

	set(pqHandshakeConfig{})
	assert.Equal(t, []bool{false, false, false}, attempts(false, false, false), "off")

	set(pqHandshakeConfig{initiate: true})
	assert.Equal(t, []bool{true, false, true, false}, attempts(false, false, false, false), "timer retries alternate, PQ first")
	assert.Equal(t, []bool{true, true, false, true, false}, attempts(false, true, false, true, false), "an early trigger is a new path, which gets PQ first")
	assert.Equal(t, []bool{true, true, false}, attempts(true, true, false), "a triggered first attempt is PQ")

	t.Run("hearing back stops the classical turns", func(t *testing.T) {
		hh := &HandshakeHostInfo{pq: &pqAttempt{}}
		hh.counter, hh.lastPQ = 1, true
		hh.counter++
		assert.False(t, hm.usePQ(hh, peer, false), "unheard, the classical attempt gets its turn")
		hh.pq.heard = true
		assert.True(t, hm.usePQ(hh, peer, false))
	})

	set(pqHandshakeConfig{initiate: true, accept: true, require: true})
	assert.Equal(t, []bool{true, true, true, true}, attempts(false, false, true, false), "require never offers classical")

	set(pqHandshakeConfig{initiate: true, sticky: true})
	assert.Equal(t, []bool{true, false}, attempts(false, false), "sticky alternates until the peer has done PQ")
	hm.notePQPeer(header.HandshakeIXPSK2MLKEM1024, []netip.Addr{peer})
	assert.Equal(t, []bool{true, true, true}, attempts(false, false, false), "then only PQ")
}