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

package canarydns

import (
	"encoding/binary"
	"hash/fnv"
	"net"
	"net/netip"
	"strings"
	"sync"
	"time"
)

// DNS constants (RFC 1035 + RFC 6891 EDNS).
const (
	typeA    = 1
	typeNS   = 2
	typeTXT  = 16
	typeAAAA = 28
	typeOPT  = 41

	classIN = 1

	rcodeNoError  = 0
	rcodeNXDomain = 3

	flagQR = 0x8000
	flagAA = 0x0400
	flagTC = 0x0200
	flagRD = 0x0100
	flagRA = 0x0080

	udpMaxNoEDNS = 512
	ednsDO       = 0x8000 // DO bit lives in the OPT TTL field's high half
	optECS       = 8      // EDNS Client Subnet option code
)

// Query is one logged canary lookup (spec §6 dns_canary shape).
type Query struct {
	QName         string    `json:"qname"`
	At            time.Time `json:"at"`
	ResolverIP    string    `json:"resolver_ip"`
	Transport     string    `json:"transport"` // "udp" | "tcp"
	EDNS          edns      `json:"edns"`
	ECS           string    `json:"ecs,omitempty"`
	CasePreserved bool      `json:"case_preserved"`
	// qname_minimized is not reliably detectable authoritative-side without
	// cross-query correlation; left false (TODO) rather than guessed.
	QNameMinimized bool `json:"qname_minimized"`
}

type edns struct {
	Present bool     `json:"present"`
	Bufsize int      `json:"bufsize"`
	Flags   []string `json:"flags"`
}

// Server is the authoritative responder for one canary zone.
type Server struct {
	zone      string // fully-qualified, lowercase, trailing dot, e.g. "c.echo-lot.app."
	nsName    string // this server's own name for NS/authority answers
	primaryV4 netip.Addr
	primaryV6 netip.Addr

	mu       sync.Mutex
	log      []Query // ring, newest last
	retainTo time.Time
}

const logCap = 8192

// New creates a server for zone (with or without trailing dot). nsName is the
// server's own hostname (for the zone's NS record); primary v4/v6 are this
// host's addresses used to answer the zone apex / NS glue.
func New(zone, nsName string, v4, v6 netip.Addr) *Server {
	z := strings.ToLower(strings.TrimSuffix(zone, ".")) + "."
	return &Server{zone: z, nsName: strings.TrimSuffix(nsName, ".") + ".", primaryV4: v4, primaryV6: v6}
}

// RecentForPrefix returns logged queries whose qname contains ".<prefix>."
// (the session prefix the app embeds: <nonce>.<session-prefix>.<zone>).
func (s *Server) RecentForPrefix(prefix string) []Query {
	s.mu.Lock()
	defer s.mu.Unlock()
	needle := "." + strings.ToLower(prefix) + "."
	var out []Query
	for _, q := range s.log {
		if strings.Contains(strings.ToLower(q.QName), needle) {
			out = append(out, q)
		}
	}
	return out
}

func (s *Server) record(q Query) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.log) >= logCap {
		s.log = s.log[1:]
	}
	s.log = append(s.log, q)
}

// ServeUDP / ServeTCP run read loops; call one per bound address.
func (s *Server) ServeUDP(conn *net.UDPConn) error {
	buf := make([]byte, 1500)
	for {
		n, raddr, err := conn.ReadFromUDPAddrPort(buf)
		if err != nil {
			return err
		}
		resp := s.handle(buf[:n], raddr.Addr(), "udp")
		if resp != nil {
			_, _ = conn.WriteToUDPAddrPort(resp, raddr)
		}
	}
}

func (s *Server) ServeTCP(ln net.Listener) error {
	for {
		c, err := ln.Accept()
		if err != nil {
			return err
		}
		go s.handleTCP(c)
	}
}

func (s *Server) handleTCP(c net.Conn) {
	defer c.Close()
	_ = c.SetDeadline(time.Now().Add(10 * time.Second))
	var lenBuf [2]byte
	if _, err := readFull(c, lenBuf[:]); err != nil {
		return
	}
	msg := make([]byte, binary.BigEndian.Uint16(lenBuf[:]))
	if _, err := readFull(c, msg); err != nil {
		return
	}
	ra, _ := netip.ParseAddrPort(c.RemoteAddr().String())
	resp := s.handle(msg, ra.Addr(), "tcp")
	if resp == nil {
		return
	}
	// TCP has no 512 limit; never truncate.
	out := make([]byte, 2+len(resp))
	binary.BigEndian.PutUint16(out[0:2], uint16(len(resp)))
	copy(out[2:], resp)
	_, _ = c.Write(out)
}

func readFull(c net.Conn, b []byte) (int, error) {
	got := 0
	for got < len(b) {
		n, err := c.Read(b[got:])
		got += n
		if err != nil {
			return got, err
		}
	}
	return got, nil
}

// handle parses one query, logs it, and returns the wire response (nil to drop).
func (s *Server) handle(pkt []byte, resolver netip.Addr, transport string) []byte {
	if len(pkt) < 12 {
		return nil
	}
	id := binary.BigEndian.Uint16(pkt[0:2])
	qdcount := binary.BigEndian.Uint16(pkt[4:6])
	arcount := binary.BigEndian.Uint16(pkt[10:12])
	if qdcount != 1 {
		return s.errorResponse(id, rcodeNoError, nil) // we only answer single-question queries
	}

	qnameRaw, qtype, _, qEnd, ok := parseQuestion(pkt, 12)
	if !ok {
		return nil
	}

	// EDNS OPT is an additional-section RR; scan for it after the question.
	opt := parseOPT(pkt, qEnd, arcount)

	// Log every query — this is the whole point of the canary zone.
	q := Query{
		QName: strings.TrimSuffix(qnameRaw, "."), At: time.Now().UTC(),
		ResolverIP: resolver.Unmap().String(), Transport: transport,
		EDNS:          opt.edns,
		ECS:           opt.ecs,
		CasePreserved: qnameRaw == strings.ToLower(qnameRaw), // mixed case ⇒ 0x20 randomization
	}
	s.record(q)

	name := strings.ToLower(qnameRaw)
	if !strings.HasSuffix(name, s.zone) {
		return s.errorResponse(id, rcodeNXDomain, &opt)
	}
	sub := strings.TrimSuffix(name, s.zone) // e.g. "ttl-5." or "" for apex

	return s.answer(id, pkt, qEnd, sub, qtype, &opt, transport)
}

// answer builds the response for a name known to be in-zone.
func (s *Server) answer(id uint16, pkt []byte, qEnd int, sub string, qtype uint16, opt *optInfo, transport string) []byte {
	labels := splitLabels(sub) // e.g. ["ttl-5"], [], ["<nonce>","miss"], ["<nonce>","<sessprefix>"]

	var rrs []rr
	switch {
	case len(labels) == 0: // zone apex
		if qtype == typeNS {
			rrs = append(rrs, rr{ttl: 3600, typ: typeNS, ns: s.nsName})
		} else if qtype == typeA && s.primaryV4.IsValid() {
			rrs = append(rrs, rr{ttl: 3600, typ: typeA, addr: s.primaryV4})
		} else if qtype == typeAAAA && s.primaryV6.IsValid() {
			rrs = append(rrs, rr{ttl: 3600, typ: typeAAAA, addr: s.primaryV6})
		}
	case len(labels) == 1:
		if ref := findReference(labels[0]); ref != nil {
			rrs = referenceAnswers(ref, qtype)
		}
	default:
		// Per-query names: <nonce>.miss.<zone> and <nonce>.<session-prefix>.<zone>.
		// Deterministic A derived from the leftmost label (the nonce), TTL 3600,
		// documentation range — ground truth that can never be pre-cached.
		if qtype == typeA {
			rrs = append(rrs, rr{ttl: 3600, typ: typeA, addr: nonceAddr(labels[0])})
		}
	}

	if len(rrs) == 0 {
		// In-zone but no such record/type → NOERROR/NODATA (or NXDOMAIN at apex miss).
		return s.buildResponse(id, pkt, qEnd, nil, opt, transport, rcodeNoError)
	}
	return s.buildResponse(id, pkt, qEnd, rrs, opt, transport, rcodeNoError)
}

// nonceAddr maps a nonce label into 192.0.2.0/24 deterministically.
func nonceAddr(nonce string) netip.Addr {
	h := fnv.New32a()
	_, _ = h.Write([]byte(nonce))
	return netip.AddrFrom4([4]byte{192, 0, 2, byte(h.Sum32()%254 + 1)})
}

func referenceAnswers(ref *refRecord, qtype uint16) []rr {
	var rrs []rr
	switch qtype {
	case typeA:
		for _, a := range ref.a {
			if a.Is4() {
				rrs = append(rrs, rr{ttl: ref.ttl, typ: typeA, addr: a})
			}
		}
	case typeAAAA:
		for _, a := range ref.a {
			if a.Is6() && !a.Is4In6() {
				rrs = append(rrs, rr{ttl: ref.ttl, typ: typeAAAA, addr: a})
			}
		}
	case typeTXT:
		if len(ref.txt) > 0 {
			rrs = append(rrs, rr{ttl: ref.ttl, typ: typeTXT, txt: ref.txt})
		}
	}
	return rrs
}

func splitLabels(sub string) []string {
	sub = strings.TrimSuffix(sub, ".")
	if sub == "" {
		return nil
	}
	return strings.Split(sub, ".")
}

func (s *Server) errorResponse(id uint16, rcode int, opt *optInfo) []byte {
	hdr := make([]byte, 12)
	binary.BigEndian.PutUint16(hdr[0:2], id)
	binary.BigEndian.PutUint16(hdr[2:4], uint16(flagQR|flagAA|rcode))
	if opt != nil && opt.edns.Present {
		binary.BigEndian.PutUint16(hdr[10:12], 1)
		return append(hdr, buildOPT(opt)...)
	}
	return hdr
}