<commit_before>package mdns

import (
	"log"
	"net"
	"strings"
	"sync"

	"github.com/miekg/dns"
)

const (
	maxMsgSize = 512
	bufSize    = 1500
)

var (
	ip4Addr   = &net.UDPAddr{IP: net.ParseIP("172.16.17.32"), Port: 5353}
	localZone = &zone{entries: make(map[string][]dns.RR)}
)

type server struct {
	conn *net.UDPConn
}

type zone struct {
	lock    sync.Mutex
	entries map[string][]dns.RR
}

func Publish(rrStr string) {
	rr, err := dns.NewRR(rrStr)
	if err != nil {
		log.Println("Failure parsing RR", err)
		return
	}
	PublishRR(rr)
}

func PublishRR(rr dns.RR) {
	hdr := rr.Header()
	if !strings.HasSuffix(hdr.Name, ".local.") {
		log.Panicln("Published domain MUST end with .local.", hdr.Name)
		return
	}

	localZone.publish(rr)
}

func init() {
	newServer(ip4Addr)
}

func newServer(addr *net.UDPAddr) *server {
	conn, err := net.ListenMulticastUDP("udp", nil, addr)
	if err != nil {
		log.Panicln("Unable to listen to MDNS port", err)
	}

	s := &server{conn: conn}
	go s.listen()
	return s
}

func (s *server) listen() {
	buf := make([]byte, bufSize)
	for {
		n, addr, err := s.conn.ReadFromUDP(buf)
		if err != nil {
			log.Panicln("Error reading UDP packet", err)
		}

		msg := &dns.Msg{}
		err = msg.Unpack(buf[:n])
		if err != nil {
			log.Println("Error unpacking DNS packet", err)
			continue
		}
		log.Println("-----------------------------------------------")
		log.Println("Addr", addr, "Packet", msg)

		var resp *dns.Msg
		if msg.MsgHdr.Response {
			resp = s.doResponse(msg)
		} else {
			resp = s.doQuestion(msg)
		}

		if resp != nil {
			b, err := resp.PackBuffer(buf)
			if err != nil {
				log.Println("Unable to pack repsonse", err)
			} else {
				n, err = s.conn.WriteToUDP(b, addr)
				if err != nil {
					log.Println("Unable to send response", err)
				}
			}
		}
	}
}

func (s *server) doQuestion(msg *dns.Msg) *dns.Msg {
	if msg.MsgHdr.Opcode != dns.OpcodeQuery {
		log.Println("Non-query opcodes not supported", msg.MsgHdr.Opcode)
		return nil
	}

	resp := new(dns.Msg)
	resp.MsgHdr = dns.MsgHdr{
		Id:            msg.MsgHdr.Id,
		Response:      true,
		Authoritative: true,
		Opcode:        dns.OpcodeQuery,
	}

	for _, q := range msg.Question {
		resp.Question = append(resp.Question, q)
		resp.Answer = append(resp.Answer, localZone.query(q)...)
	}

	return resp
}

func (s *server) doResponse(msg *dns.Msg) *dns.Msg {
	return nil
}

func (z *zone) query(q dns.Question) []dns.RR {
	z.lock.Lock()
	defer z.lock.Unlock()

	name := q.Name
	entries := z.entries[name]
<commit_msg>Implement query.<commit_after>package mdns

import (
	"log"
	"net"
	"strings"
	"sync"

	"github.com/miekg/dns"
)

const (
	maxMsgSize = 512
	bufSize    = 1500
)

var (
	ip4Addr   = &net.UDPAddr{IP: net.ParseIP("172.16.17.32"), Port: 5353}
	localZone = &zone{entries: make(map[string][]dns.RR)}

	ip4Server *server
)

type Response struct {
	Chan <-chan dns.RR

	ch   chan<- dns.RR
	done chan bool
	q    dns.Question
}

type server struct {
	conn *net.UDPConn

	lock    sync.Mutex
	queries []*Response
}

type zone struct {
	lock    sync.Mutex
	entries map[string][]dns.RR
}

func Publish(rrStr string) {
	rr, err := dns.NewRR(rrStr)
	if err != nil {
		log.Println("Failure parsing RR", err)
		return
	}
	PublishRR(rr)
}

func PublishRR(rr dns.RR) {
	hdr := rr.Header()
	if !strings.HasSuffix(hdr.Name, ".local.") {
		log.Panicln("Published domain MUST end with .local.", hdr.Name)
		return
	}

	localZone.publish(rr)
}

func Query(name string) *Response {
	ch := make(chan dns.RR)
	resp := &Response{Chan: ch, ch: ch, done: make(chan bool)}
	resp.q = dns.Question{Name: name, Qtype: dns.TypeANY, Qclass: dns.ClassINET}
	go resp.query()
	return resp
}

func (r *Response) Done() {
	close(r.done)
}

func (r *Response) query() {
	ip4Server.query(r)

	<-r.done
	ip4Server.endQuery(r)
}

func (r *Response) answer(rr dns.RR) {
	select {
	case <-r.done:
		// Do nothing, query is finished.
	case r.ch <- rr:
		// Answer delivered.
	}
}

func init() {
	ip4Server = newServer(ip4Addr)
}

func newServer(addr *net.UDPAddr) *server {
	conn, err := net.ListenMulticastUDP("udp", nil, addr)
	if err != nil {
		log.Panicln("Unable to listen to MDNS port", err)
	}

	s := &server{conn: conn}
	go s.listen()
	return s
}

func (s *server) query(r *Response) {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.queries = append(s.queries, r)

	msg := new(dns.Msg)
	// TODO: Use ID
	// msg.Id = dns.Id()
	msg.Question = []dns.Question{r.q}
	err := s.send(msg, s.conn.LocalAddr().(*net.UDPAddr))
	if err != nil {
		log.Println("Error sending query", err)
	}
}

func (s *server) endQuery(r *Response) {
	s.lock.Lock()
	defer s.lock.Unlock()

	for i, q := range s.queries {
		if q == r {
			last := len(s.queries) - 1
			s.queries[i] = s.queries[last]
			s.queries[last] = nil
			s.queries = s.queries[:last]
			break
		}
	}
}

func (s *server) send(msg *dns.Msg, addr *net.UDPAddr) error {
	b, err := msg.Pack()
	if err != nil {
		return err
	}

	_, err = s.conn.WriteToUDP(b, addr)
	if err != nil {
		return err
	}
	return nil
}

func (s *server) listen() {
	buf := make([]byte, bufSize)
	for {
		n, addr, err := s.conn.ReadFromUDP(buf)
		if err != nil {
			log.Panicln("Error reading UDP packet", err)
		}

		msg := &dns.Msg{}
		err = msg.Unpack(buf[:n])
		if err != nil {
			log.Println("Error unpacking DNS packet", err)
			continue
		}
		log.Println("-----------------------------------------------")
		log.Println("Addr", addr, "Packet", msg)

		var resp *dns.Msg
		if msg.MsgHdr.Response {
			resp = s.doResponse(msg)
		} else {
			resp = s.doQuestion(msg)
		}

		if resp != nil {
			err = s.send(resp, addr)
			if err != nil {
				log.Println("Unable to send response", err)
			}
		}
	}
}

func (s *server) doQuestion(msg *dns.Msg) *dns.Msg {
	if msg.MsgHdr.Opcode != dns.OpcodeQuery {
		log.Println("Non-query opcodes not supported", msg.MsgHdr.Opcode)
		return nil
	}

	resp := new(dns.Msg)
	resp.MsgHdr = dns.MsgHdr{
		Id:            msg.MsgHdr.Id,
		Response:      true,
		Authoritative: true,
		Opcode:        dns.OpcodeQuery,
	}

	for _, q := range msg.Question {
		resp.Question = append(resp.Question, q)
		resp.Answer = append(resp.Answer, localZone.query(q)...)
	}

	return resp
}

func (s *server) doResponse(msg *dns.Msg) *dns.Msg {
	if len(msg.Question) != 1 {
		// TODO: Support requests/responses with multiple questions
		return nil
	}

	s.lock.Lock()
	defer s.lock.Unlock()

	for _, q := range s.queries {
		if q.q == msg.Question[0] {
			for _, a := range msg.Answer {
				q.answer(a)
			}
		}
	}

	return nil
}

func (z *zone) query(q dns.Question) []dns.RR {
	z.lock.Lock()
	defer z.lock.Unlock()

	name := q.Name
	entries := z.entries[name]

{'fraction_non_alphanumeric': 0.17537537537537537, 'fraction_numerical': 0.0076576576576576575, 'mean_word_length': 2.897600936220012, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, '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': '2790871', 'n_tokens_mistral': 2912, 'n_tokens_neox': 2699, 'n_words': 635}