minidns.go

package main

import (
	"flag"
	"net"
	"strings"

	"github.com/miekg/dns"
)

type dnsHandler struct {
	table map[string][]string
}

func (d *dnsHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)
	answerer := d.newAnswer(r.Question[0].Qtype)
	answerer(m)
	w.WriteMsg(m)
}

func (d *dnsHandler) newAnswer(qType uint16) func(m *dns.Msg) {
	switch qType {
	case dns.TypeA:
		return d.answerTypeA
	}
	return d.answerNothing
}

func (d *dnsHandler) answerNothing(_ *dns.Msg) {}

func (d *dnsHandler) answerTypeA(m *dns.Msg) {
	m.Authoritative = true
	domain := m.Question[0].Name
	addresses, ok := d.table[domain]
	if !ok {
		return
	}
	for _, address := range addresses {
		m.Answer = append(m.Answer, &dns.A{
			Hdr: dns.RR_Header{Name: domain, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 10},
			A:   net.ParseIP(address),
		})
	}
}

func NewDnsHandler(configs string) dns.Handler {
	table := make(map[string][]string)
	for _, config := range strings.Split(configs, ";") {
		items := strings.Split(config, ":")
		if len(items) != 2 {
			continue
		}
		domainName, values := items[0], items[1]
		ipAddress := strings.Split(values, ",")
		table[domainName] = ipAddress
	}
	return &dnsHandler{table: table}
}

func main() {
	var configs string
	var addr string
	flag.StringVar(&addr, "addr", "0.0.0.0:53", "address to listen")
	flag.StringVar(&configs, "configs", "", "for example, example.com:192.168.0.100,192.168.0.200;example2.com:192.168.1.100")
	flag.Parse()

	handler := NewDnsHandler(configs)
	if err := dns.ListenAndServe(addr, "udp", handler); err != nil {
		panic(err)
	}
}