package main

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"math/rand/v2"
	"net"
	"strings"
	"syscall"
	"time"
)

// A real question to the resolver, as any program on the machine asks it: one name, one record type, over
// UDP or TCP, to one address. Written out here rather than taken from a library: the module needs the
// header's answer code and the A records, nothing else, and Go's own resolver would read /etc/hosts and
// /etc/resolv.conf instead of asking the one address it is told (novox/hq issue 348).

// Answer is what one question got.
type Answer struct {
	// Server is where it was asked, and Over the transport.
	Server string `json:"server"`
	Over   string `json:"over"`
	Name   string `json:"name"`
	// Answered says a reply came; Code is its answer code (NOERROR, NXDOMAIN, …) and Addresses its A
	// records.
	Answered  bool     `json:"answered"`
	Code      string   `json:"code,omitempty"`
	Addresses []string `json:"addresses,omitempty"`
	// Failed is why no reply came, in words: refused (nothing listens there), no answer within, or what
	// the system said.
	Failed string `json:"failed,omitempty"`
	// Took is how long it took, in milliseconds.
	Took int64 `json:"took_ms"`
}

// rcodes are the answer codes by number.
var rcodes = map[int]string{0: "NOERROR", 1: "FORMERR", 2: "SERVFAIL", 3: "NXDOMAIN", 4: "NOTIMP", 5: "REFUSED"}

// query is one question for name's A record, with id.
func query(name string, id uint16) ([]byte, error) {
	name = strings.TrimSuffix(name, ".")
	if name == "" {
		return nil, errors.New("no name to ask")
	}
	msg := make([]byte, 12, 64)
	binary.BigEndian.PutUint16(msg[0:], id)
	binary.BigEndian.PutUint16(msg[2:], 0x0100) // a standard query, recursion desired
	binary.BigEndian.PutUint16(msg[4:], 1)      // one question
	for _, label := range strings.Split(name, ".") {
		if label == "" || len(label) > 63 {
			return nil, fmt.Errorf("%q is not a name that can be asked", name)
		}
		msg = append(msg, byte(len(label)))
		msg = append(msg, label...)
	}
	msg = append(msg, 0, 0, 1, 0, 1) // root; type A; class IN
	return msg, nil
}

// readAnswer reads a reply to the question with id: its code and A records.
func readAnswer(msg []byte, id uint16) (string, []string, error) {
	if len(msg) < 12 {
		return "", nil, errors.New("a reply too short to be one")
	}
	if binary.BigEndian.Uint16(msg[0:]) != id {
		return "", nil, errors.New("a reply to another question")
	}
	flags := binary.BigEndian.Uint16(msg[2:])
	code, ok := rcodes[int(flags&0xF)]
	if !ok {
		code = fmt.Sprintf("RCODE%d", flags&0xF)
	}
	qd, an := int(binary.BigEndian.Uint16(msg[4:])), int(binary.BigEndian.Uint16(msg[6:]))
	at := 12
	for i := 0; i < qd; i++ {
		var err error
		if at, err = skipName(msg, at); err != nil {
			return code, nil, err
		}
		at += 4
	}
	var addresses []string
	for i := 0; i < an; i++ {
		var err error
		if at, err = skipName(msg, at); err != nil {
			return code, addresses, err
		}
		if at+10 > len(msg) {
			return code, addresses, errors.New("a reply cut short")
		}
		kind, length := binary.BigEndian.Uint16(msg[at:]), int(binary.BigEndian.Uint16(msg[at+8:]))
		at += 10
		if at+length > len(msg) {
			return code, addresses, errors.New("a reply cut short")
		}
		if kind == 1 && length == 4 {
			addresses = append(addresses, net.IP(msg[at:at+4]).String())
		}
		at += length
	}
	return code, addresses, nil
}

// skipName steps over a name in a message, a pointer ending it.
func skipName(msg []byte, at int) (int, error) {
	for at < len(msg) {
		n := int(msg[at])
		switch {
		case n == 0:
			return at + 1, nil
		case n&0xC0 == 0xC0:
			return at + 2, nil
		}
		at += n + 1
	}
	return at, errors.New("a name runs past the reply")
}

// Ask asks server (an address, port 53 unless given) for name's A record over "udp" or "tcp", within
// timeout.
func Ask(ctx context.Context, server, over, name string, timeout time.Duration) Answer {
	if _, _, err := net.SplitHostPort(server); err != nil {
		server = net.JoinHostPort(server, "53")
	}
	a := Answer{Server: server, Over: over, Name: name}
	started := time.Now()
	defer func() { a.Took = time.Since(started).Milliseconds() }()
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	id := uint16(rand.UintN(1 << 16))
	q, err := query(name, id)
	if err != nil {
		a.Failed = err.Error()
		return a
	}
	conn, err := (&net.Dialer{}).DialContext(ctx, over, server)
	if err != nil {
		a.Failed = failedWords(err, timeout)
		a.Took = time.Since(started).Milliseconds()
		return a
	}
	defer conn.Close()
	if deadline, ok := ctx.Deadline(); ok {
		_ = conn.SetDeadline(deadline)
	}
	var reply []byte
	if over == "tcp" {
		framed := binary.BigEndian.AppendUint16(nil, uint16(len(q)))
		if _, err = conn.Write(append(framed, q...)); err == nil {
			head := make([]byte, 2)
			if _, err = readFull(conn, head); err == nil {
				reply = make([]byte, binary.BigEndian.Uint16(head))
				_, err = readFull(conn, reply)
			}
		}
	} else if _, err = conn.Write(q); err == nil {
		buf := make([]byte, 1500)
		var n int
		n, err = conn.Read(buf)
		reply = buf[:n]
	}
	if err != nil {
		a.Failed = failedWords(err, timeout)
		a.Took = time.Since(started).Milliseconds()
		return a
	}
	a.Answered = true
	a.Code, a.Addresses, err = readAnswer(reply, id)
	if err != nil {
		a.Failed = err.Error()
	}
	a.Took = time.Since(started).Milliseconds()
	return a
}

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

// failedWords is why no reply came: refused when nothing listens, no answer when nothing came in time.
func failedWords(err error, timeout time.Duration) string {
	var timedOut interface{ Timeout() bool }
	switch {
	case errors.Is(err, syscall.ECONNREFUSED):
		return "refused: nothing listens there"
	case errors.Is(err, context.DeadlineExceeded) || (errors.As(err, &timedOut) && timedOut.Timeout()):
		return fmt.Sprintf("no answer within %s", timeout)
	}
	return err.Error()
}