package net

import (
	"context"
	"github.com/benburkert/dns"
	"net"
)

func ReplaceDefaultDnsResolver(){
	cli := &dns.Client{Resolver: &PreferIPV4Resolver{}}
	net.DefaultResolver.PreferGo = true
	net.DefaultResolver.Dial = cli.Dial
}

type PreferIPV4Resolver struct{}

func (p *PreferIPV4Resolver) ServeDNS(ctx context.Context, w dns.MessageWriter, r *dns.Query) {
	if r.Questions[0].Type == dns.TypeAAAA {
		w.Status(dns.ServFail)
	} else if msg, err := w.Recur(ctx); err != nil {
		w.Status(dns.ServFail)
	} else {
		p.writeMessage(w, msg)
	}
}

func (p *PreferIPV4Resolver) writeMessage(w dns.MessageWriter, msg *dns.Message) {
	w.Status(msg.RCode)
	w.Authoritative(msg.Authoritative)
	w.Recursion(msg.RecursionAvailable)

	for _, res := range msg.Answers {
		w.Answer(res.Name, res.TTL, res.Record)
	}
	for _, res := range msg.Authorities {
		w.Authority(res.Name, res.TTL, res.Record)
	}
	for _, res := range msg.Additionals {
		w.Additional(res.Name, res.TTL, res.Record)
	}
}