<commit_before>package server

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/miekg/dns"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/dice"
	"v2ray.com/core/common/net"
	"v2ray.com/core/transport/internet/udp"
)

const (
	CleanupInterval  = time.Second * 120
	CleanupThreshold = 512
)

var (
	pseudoDestination = net.UDPDestination(net.LocalHostIP, net.Port(53))
)

type ARecord struct {
	IPs    []net.IP
	Expire time.Time
}

type NameServer interface {
	QueryA(domain string) <-chan *ARecord
}

type PendingRequest struct {
	expire   time.Time
	response chan<- *ARecord
}

type UDPNameServer struct {
	sync.Mutex
	address     net.Destination
	requests    map[uint16]*PendingRequest
	udpServer   *udp.Dispatcher
	nextCleanup time.Time
}

func NewUDPNameServer(address net.Destination, dispatcher dispatcher.Interface) *UDPNameServer {
	s := &UDPNameServer{
		address:   address,
		requests:  make(map[uint16]*PendingRequest),
		udpServer: udp.NewDispatcher(dispatcher),
	}
	return s
}

func (v *UDPNameServer) Cleanup() {
	expiredRequests := make([]uint16, 0, 16)
	now := time.Now()
	v.Lock()
	for id, r := range v.requests {
		if r.expire.Before(now) {
			expiredRequests = append(expiredRequests, id)
			close(r.response)
		}
	}
	for _, id := range expiredRequests {
		delete(v.requests, id)
	}
	v.Unlock()
}

func (v *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
	var id uint16
	v.Lock()
	if len(v.requests) > CleanupThreshold && v.nextCleanup.Before(time.Now()) {
		v.nextCleanup = time.Now().Add(CleanupInterval)
		go v.Cleanup()
	}

	for {
		id = dice.RollUint16()
		if _, found := v.requests[id]; found {
			continue
		}
		log.Trace(newError("add pending request id ", id).AtDebug())
		v.requests[id] = &PendingRequest{
			expire:   time.Now().Add(time.Second * 8),
			response: response,
		}
		break
	}
	v.Unlock()
	return id
}

func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) {
	msg := new(dns.Msg)
	err := msg.Unpack(payload.Bytes())
	if err != nil {
		log.Trace(newError("failed to parse DNS response").Base(err).AtWarning())
		return
	}
	record := &ARecord{
		IPs: make([]net.IP, 0, 16),
	}
	id := msg.Id
	ttl := uint32(3600) // an hour
	log.Trace(newError("handling response for id ", id, " content: ", msg).AtDebug())

	v.Lock()
	request, found := v.requests[id]
	if !found {
		v.Unlock()
		return
	}
	delete(v.requests, id)
	v.Unlock()

	for _, rr := range msg.Answer {
		switch rr := rr.(type) {
		case *dns.A:
			record.IPs = append(record.IPs, rr.A)
			fmt.Println("Adding ans:", rr.A)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		case *dns.AAAA:
			record.IPs = append(record.IPs, rr.AAAA)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		}
	}
	record.Expire = time.Now().Add(time.Second * time.Duration(ttl))

	request.response <- record
	close(request.response)
}

func (v *UDPNameServer) BuildQueryA(domain string, id uint16) *buf.Buffer {

	msg := new(dns.Msg)
	msg.Id = id
	msg.RecursionDesired = true
	msg.Question = []dns.Question{<commit_msg>remove debug info<commit_after>package server

import (
	"context"
	"sync"
	"time"

	"github.com/miekg/dns"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/dice"
	"v2ray.com/core/common/net"
	"v2ray.com/core/transport/internet/udp"
)

const (
	CleanupInterval  = time.Second * 120
	CleanupThreshold = 512
)

var (
	pseudoDestination = net.UDPDestination(net.LocalHostIP, net.Port(53))
)

type ARecord struct {
	IPs    []net.IP
	Expire time.Time
}

type NameServer interface {
	QueryA(domain string) <-chan *ARecord
}

type PendingRequest struct {
	expire   time.Time
	response chan<- *ARecord
}

type UDPNameServer struct {
	sync.Mutex
	address     net.Destination
	requests    map[uint16]*PendingRequest
	udpServer   *udp.Dispatcher
	nextCleanup time.Time
}

func NewUDPNameServer(address net.Destination, dispatcher dispatcher.Interface) *UDPNameServer {
	s := &UDPNameServer{
		address:   address,
		requests:  make(map[uint16]*PendingRequest),
		udpServer: udp.NewDispatcher(dispatcher),
	}
	return s
}

func (v *UDPNameServer) Cleanup() {
	expiredRequests := make([]uint16, 0, 16)
	now := time.Now()
	v.Lock()
	for id, r := range v.requests {
		if r.expire.Before(now) {
			expiredRequests = append(expiredRequests, id)
			close(r.response)
		}
	}
	for _, id := range expiredRequests {
		delete(v.requests, id)
	}
	v.Unlock()
}

func (v *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
	var id uint16
	v.Lock()
	if len(v.requests) > CleanupThreshold && v.nextCleanup.Before(time.Now()) {
		v.nextCleanup = time.Now().Add(CleanupInterval)
		go v.Cleanup()
	}

	for {
		id = dice.RollUint16()
		if _, found := v.requests[id]; found {
			continue
		}
		log.Trace(newError("add pending request id ", id).AtDebug())
		v.requests[id] = &PendingRequest{
			expire:   time.Now().Add(time.Second * 8),
			response: response,
		}
		break
	}
	v.Unlock()
	return id
}

func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) {
	msg := new(dns.Msg)
	err := msg.Unpack(payload.Bytes())
	if err != nil {
		log.Trace(newError("failed to parse DNS response").Base(err).AtWarning())
		return
	}
	record := &ARecord{
		IPs: make([]net.IP, 0, 16),
	}
	id := msg.Id
	ttl := uint32(3600) // an hour
	log.Trace(newError("handling response for id ", id, " content: ", msg).AtDebug())

	v.Lock()
	request, found := v.requests[id]
	if !found {
		v.Unlock()
		return
	}
	delete(v.requests, id)
	v.Unlock()

	for _, rr := range msg.Answer {
		switch rr := rr.(type) {
		case *dns.A:
			record.IPs = append(record.IPs, rr.A)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		case *dns.AAAA:
			record.IPs = append(record.IPs, rr.AAAA)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		}
	}
	record.Expire = time.Now().Add(time.Second * time.Duration(ttl))

	request.response <- record
	close(request.response)
}

func (v *UDPNameServer) BuildQueryA(domain string, id uint16) *buf.Buffer {

	msg := new(dns.Msg)
	msg.Id = id
	msg.RecursionDesired = true
	msg.Question = []dns.Question{
{'fraction_non_alphanumeric': 0.15151025682442254, 'fraction_numerical': 0.013568082700694557, 'mean_word_length': 3.666164280331575, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '332296', 'n_tokens_mistral': 2650, 'n_tokens_neox': 2476, 'n_words': 491}