cool-dns/coolDns.go

330 lines
6.9 KiB
Go

package main
import (
"io/ioutil"
"log"
"net"
"os"
"os/signal"
"strconv"
"strings"
"syscall"
"github.com/miekg/dns"
"gopkg.in/yaml.v3"
)
type zone struct {
zone string
rr rrMap
acl []string
}
type rrMap map[uint16]map[string][]dns.RR
type config struct {
Zones []configZone `yaml:"zones"`
ACL []configACL `yaml:"acl"`
Forward configForward `yaml:"forward"`
}
type configForward struct {
ACL []string `yaml:"acl"`
Server string `yaml:"server"`
}
type configACL struct {
Name string `yaml:"name"`
CIDR string `yaml:"cidr"`
}
type configZone struct {
Zone string `yaml:"zone"`
File string `yaml:"file"`
ACL []string `yaml:"acl"`
}
var anyRecordTypes = []uint16{
dns.TypeSOA,
dns.TypeA,
dns.TypeAAAA,
dns.TypeNS,
dns.TypeCNAME,
dns.TypeMX,
dns.TypeTXT,
dns.TypeSRV,
dns.TypeCAA,
}
func loadConfig() (*config, error) {
file, err := ioutil.ReadFile("config.yml")
if err != nil {
return nil, err
}
var loadedConfig config
err = yaml.Unmarshal(file, &loadedConfig)
if err != nil {
return nil, err
}
return &loadedConfig, nil
}
func loadZones(configZones []configZone) ([]zone, error) {
zones := make([]zone, 0)
for _, z := range configZones {
rrs, err := loadZonefile(z.File)
if err != nil {
return nil, err
}
zones = append(zones, zone{
zone: z.Zone,
rr: createRRMap(rrs),
acl: z.ACL,
})
log.Printf("Loaded zone %s\n", z.Zone)
}
return zones, nil
}
func createRRMap(rrs []dns.RR) rrMap {
rrMap := make(rrMap)
for _, rr := range rrs {
if rrMap[rr.Header().Rrtype] == nil {
rrMap[rr.Header().Rrtype] = make(map[string][]dns.RR)
}
if rrMap[rr.Header().Rrtype][rr.Header().Name] == nil {
rrMap[rr.Header().Rrtype][rr.Header().Name] = make([]dns.RR, 0)
}
rrMap[rr.Header().Rrtype][rr.Header().Name] = append(rrMap[rr.Header().Rrtype][rr.Header().Name], rr)
}
return rrMap
}
func loadZonefile(filepath string) ([]dns.RR, error) {
file, err := os.Open(filepath)
if err != nil {
return nil, err
}
parser := dns.NewZoneParser(file, "", "")
var rrs = make([]dns.RR, 0)
for rr, ok := parser.Next(); ok; rr, ok = parser.Next() {
rrs = append(rrs, rr)
}
if err := parser.Err(); err != nil {
log.Println(err)
}
return rrs, nil
}
func createACLList(config []configACL) (map[string]*net.IPNet, error) {
acls := make(map[string]*net.IPNet)
for _, aclRule := range config {
_, mask, err := net.ParseCIDR(aclRule.CIDR)
if err != nil {
return nil, err
}
acls[aclRule.Name] = mask
}
return acls, nil
}
func createServer(zones []zone, config config, aclList map[string]*net.IPNet) *dns.ServeMux {
srv := dns.NewServeMux()
c := new(dns.Client)
for _, z := range zones {
srv.HandleFunc(z.zone, func(w dns.ResponseWriter, r *dns.Msg) {
// Parse IP
remoteIP, _, err := net.SplitHostPort(w.RemoteAddr().String())
ip := net.ParseIP(remoteIP)
if err != nil && ip != nil {
log.Printf("Faild to parse remote IP WTF? :%s", err.Error())
return
}
// Check ACL rules
if !checkACL(z.acl, aclList, ip) {
rcodeRequest(w, r, dns.RcodeRefused)
return
}
m := new(dns.Msg)
m.SetReply(r)
m.Authoritative = true
// maybe only support one question per query like most servers do it ???
for _, q := range r.Question {
rrs := z.rr[q.Qtype]
// Handle ANY
if q.Qtype == dns.TypeANY {
for _, rrType := range anyRecordTypes {
m.Answer = append(m.Answer, z.rr[rrType][q.Name]...)
}
} else {
// Handle any other type
m.Answer = append(m.Answer, rrs[q.Name]...)
// Check for wildcard
if len(m.Answer) == 0 {
parts := dns.SplitDomainName(q.Name)[1:]
searchDomain := "*." + dns.Fqdn(strings.Join(parts, "."))
foundDomain := rrs[searchDomain]
for _, rr := range foundDomain {
newRR := rr
newRR.Header().Name = q.Name
m.Answer = append(m.Answer, newRR)
}
}
}
// Handle extras
switch q.Qtype {
// Dont handle extra stuff when answering ANY request
// case dns.TypeANY:
// fallthrough
case dns.TypeMX:
// Resolve MX domains
for _, mxRR := range m.Answer {
if t, ok := mxRR.(*dns.MX); ok {
m.Extra = append(m.Extra, z.rr[dns.TypeA][t.Mx]...)
m.Extra = append(m.Extra, z.rr[dns.TypeAAAA][t.Mx]...)
}
}
case dns.TypeA, dns.TypeAAAA:
if len(m.Answer) == 0 {
// no A or AAAA found. Look for CNAME
m.Answer = append(m.Answer, z.rr[dns.TypeCNAME][q.Name]...)
if len(m.Answer) != 0 {
// Resolve CNAME
for _, nameRR := range m.Answer {
if t, ok := nameRR.(*dns.CNAME); ok {
m.Answer = append(m.Answer, z.rr[q.Qtype][t.Target]...)
}
}
}
}
case dns.TypeNS:
// Resove NS records
for _, nsRR := range m.Answer {
if t, ok := nsRR.(*dns.NS); ok {
m.Extra = append(m.Extra, z.rr[dns.TypeA][t.Ns]...)
m.Extra = append(m.Extra, z.rr[dns.TypeAAAA][t.Ns]...)
}
}
}
}
w.WriteMsg(m)
})
}
// Handle any other request
srv.HandleFunc(".", func(w dns.ResponseWriter, r *dns.Msg) {
remoteIP, _, err := net.SplitHostPort(w.RemoteAddr().String())
ip := net.ParseIP(remoteIP)
if err != nil && ip != nil {
log.Printf("Faild to parse remote IP WTF? :%s", err.Error())
return
}
// Check ACL rules
if !checkACL(config.Forward.ACL, aclList, ip) {
rcodeRequest(w, r, dns.RcodeRefused)
return
}
// Forward request
in, _, err := c.Exchange(r, config.Forward.Server)
if err != nil {
rcodeRequest(w, r, dns.RcodeServerFailure)
return
}
w.WriteMsg(in)
})
return srv
}
func listenAndServer(server *dns.ServeMux) {
go func() {
if err := dns.ListenAndServe(":"+strconv.Itoa(8053), "udp", server); err != nil {
log.Fatalf("Failed to set udp listener %s\n", err.Error())
}
}()
go func() {
if err := dns.ListenAndServe(":"+strconv.Itoa(8053), "tcp", server); err != nil {
log.Fatalf("Failed to set tcp listener %s\n", err.Error())
}
}()
sig := make(chan os.Signal)
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
s := <-sig
log.Printf("Signal (%v) received, stopping\n", s)
os.Exit(0)
}
func checkACL(alcRules []string, aclList map[string]*net.IPNet, ip net.IP) bool {
if len(alcRules) != 0 {
passed := false
for _, rule := range alcRules {
if aclList[rule].Contains(ip) {
passed = true
}
}
return passed
}
return true
}
func rcodeRequest(w dns.ResponseWriter, r *dns.Msg, rcode int) {
m := new(dns.Msg)
m.SetReply(r)
m.SetRcode(r, rcode)
w.WriteMsg(m)
}
func main() {
config, err := loadConfig()
if err != nil {
log.Fatalf("Failed to load config: %s", err.Error())
}
zones, err := loadZones(config.Zones)
if err != nil {
log.Fatalf("Failed to load zones: %s", err.Error())
}
aclList, err := createACLList(config.ACL)
if err != nil {
log.Fatalf("Failed to parse ACL rules: %s", err.Error())
}
server := createServer(zones, *config, aclList)
listenAndServer(server)
}