// SPDX-FileCopyrightText: 2026 Echolot contributors
// SPDX-License-Identifier: GPL-3.0-or-later

package dataplane

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"net"
	"net/netip"
	"testing"
	"time"

	"echo-lot.app/server/internal/session"
)

// craft builds a spec-§3.1 packet for a session the way the Android client will.
func craft(t *testing.T, sess *session.Session, typ byte, seq uint32, payload []byte) []byte {
	t.Helper()
	pkt := make([]byte, HeaderSize+len(payload))
	copy(pkt[0:4], Magic)
	pkt[4] = typ
	binary.BigEndian.PutUint16(pkt[6:8], uint16(len(payload)))
	prefix, err := hex.DecodeString(sess.ID[:16])
	if err != nil {
		t.Fatal(err)
	}
	copy(pkt[8:16], prefix)
	binary.BigEndian.PutUint32(pkt[16:20], seq)
	binary.BigEndian.PutUint64(pkt[20:28], uint64(time.Now().UnixNano()))
	copy(pkt[HeaderSize:], payload)
	mac := hmac.New(sha256.New, sess.Key[:])
	mac.Write(pkt[0:28])
	mac.Write(payload)
	copy(pkt[28:32], mac.Sum(nil)[:4])
	return pkt
}

func startServer(t *testing.T) (*session.Manager, netip.AddrPort) {
	t.Helper()
	mgr := session.NewManager(time.Minute)
	conn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { conn.Close() })
	go (&Server{Sessions: mgr}).Serve(conn)
	return mgr, conn.LocalAddr().(*net.UDPAddr).AddrPort()
}

func TestEchoRoundtripObservationAndAntiAmplification(t *testing.T) {
	mgr, addr := startServer(t)
	sess, _, err := mgr.New("dev1", "credential-ikm", netip.MustParseAddr("127.0.0.1"))
	if err != nil {
		t.Fatal(err)
	}

	client, err := net.DialUDP("udp", nil, net.UDPAddrFromAddrPort(addr))
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	client.SetDeadline(time.Now().Add(2 * time.Second))

	padding := make([]byte, 64) // request: 32 hdr + 64 payload = 96 bytes
	req := craft(t, sess, TypeEchoReq, 1, padding)
	if _, err := client.Write(req); err != nil {
		t.Fatal(err)
	}
	resp := make([]byte, 1500)
	n, err := client.Read(resp)
	if err != nil {
		t.Fatalf("no ECHO_RESP: %v", err)
	}
	if n > len(req) {
		t.Fatalf("anti-amplification violated: resp %d > req %d", n, len(req))
	}
	if resp[4] != TypeEchoResp {
		t.Fatalf("type = %#x, want ECHO_RESP", resp[4])
	}
	// Verify the server's HMAC with the same derived key.
	mac := hmac.New(sha256.New, sess.Key[:])
	mac.Write(resp[0:28])
	mac.Write(resp[HeaderSize:n])
	if !hmac.Equal(mac.Sum(nil)[:4], resp[28:32]) {
		t.Fatal("response HMAC does not verify")
	}
	// Observation block: observed source port must match our socket.
	obs := resp[HeaderSize:n]
	if len(obs) < 40 {
		t.Fatalf("observation block truncated to %d (req was big enough for 40)", len(obs))
	}
	gotPort := binary.BigEndian.Uint16(obs[32:34])
	wantPort := client.LocalAddr().(*net.UDPAddr).AddrPort().Port()
	if gotPort != wantPort {
		t.Fatalf("observed port %d, want %d", gotPort, wantPort)
	}
	if rcvd := binary.BigEndian.Uint32(obs[36:40]); rcvd != uint32(len(req)) {
		t.Fatalf("observed size %d, want %d", rcvd, len(req))
	}
}

func TestMtuProbeAckReportsReceivedSizeAndDoesNotAmplify(t *testing.T) {
	mgr, addr := startServer(t)
	sess, _, err := mgr.New("dev1", "credential-ikm", netip.MustParseAddr("127.0.0.1"))
	if err != nil {
		t.Fatal(err)
	}
	client, err := net.DialUDP("udp", nil, net.UDPAddrFromAddrPort(addr))
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()
	client.SetDeadline(time.Now().Add(2 * time.Second))

	// A large probe: 32 header + 1400 payload.
	probe := craft(t, sess, TypeMtuProbe, 1, make([]byte, 1400))
	if _, err := client.Write(probe); err != nil {
		t.Fatal(err)
	}
	buf := make([]byte, 2000)
	n, err := client.Read(buf)
	if err != nil {
		t.Fatalf("no MTU_ACK: %v", err)
	}
	if buf[4] != TypeMtuAck {
		t.Fatalf("type = %#x, want MTU_ACK", buf[4])
	}
	if n >= len(probe) {
		t.Fatalf("MTU_ACK (%d) must be far smaller than the probe (%d)", n, len(probe))
	}
	if got := binary.BigEndian.Uint32(buf[HeaderSize:n]); int(got) != len(probe) {
		t.Fatalf("acked size %d, want %d", got, len(probe))
	}
}

func TestDropsReplayBadHmacAndUnknownPrefix(t *testing.T) {
	mgr, addr := startServer(t)
	sess, _, err := mgr.New("dev1", "credential-ikm", netip.MustParseAddr("127.0.0.1"))
	if err != nil {
		t.Fatal(err)
	}
	client, err := net.DialUDP("udp", nil, net.UDPAddrFromAddrPort(addr))
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()

	// Valid packet once...
	req := craft(t, sess, TypeEchoReq, 7, make([]byte, 48))
	client.SetDeadline(time.Now().Add(2 * time.Second))
	client.Write(req)
	buf := make([]byte, 1500)
	if _, err := client.Read(buf); err != nil {
		t.Fatalf("first send should get a response: %v", err)
	}

	drop := func(name string, pkt []byte) {
		t.Helper()
		client.SetDeadline(time.Now().Add(300 * time.Millisecond))
		client.Write(pkt)
		if _, err := client.Read(buf); err == nil {
			t.Fatalf("%s: got a response, want silent drop", name)
		}
	}
	// ...replayed seq: silence.
	drop("replay", req)
	// Bad HMAC: silence.
	bad := craft(t, sess, TypeEchoReq, 8, make([]byte, 48))
	bad[31] ^= 0xFF
	drop("bad hmac", bad)
	// Unknown session prefix: silence.
	unk := craft(t, sess, TypeEchoReq, 9, make([]byte, 48))
	unk[8] ^= 0xFF
	drop("unknown prefix", unk)
}