mserver.go

package mserver

import (
	"fmt"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"context"

	"github.com/fsnotify/fsnotify"
	"github.com/go-redis/redis/v8"
	"github.com/miekg/dns"
	"github.com/patrickmn/go-cache"
)

var serverurl string
var c *cache.Cache
var question *cache.Cache
var redisdb *redis.Client
var ctx context.Context

var resolvconf string

// For blocklist
var blocklist map[string]bool

// We will read and fill up our blocklist
func fillBlockList() {
	blocklist = make(map[string]bool)
	var filenames []string
	// Find all the files in the blocklists directory
	filepath.Walk("blocklists", func(path string, info os.FileInfo, err error) error {

		// Need files ready for simple loading to our tool
		if !info.IsDir() && strings.Contains(path, "clean-") && strings.HasSuffix(path, ".txt") {

			filenames = append(filenames, path)
		}
		return nil
	})

	for _, filename := range filenames {
		data, err := os.ReadFile(filename)
		if err == nil {
			lines := strings.Split(string(data), "\n")
			for _, line := range lines {
				if len(line) > 0 {
					blocklist[line+"."] = true
				}
			}
			fmt.Println("Loaded:", filename)
		} else {
			fmt.Println("Error reading blocklist")
		}

	}
}

// This function monitors the /etc/resolv.conf for any change
func watchDNSConfiguration(watcher fsnotify.Watcher) {
	for {
		select {
		case event, ok := <-watcher.Events:
			if !ok {
				return
			}
			//log.Println("event:", event)
			if event.Has(fsnotify.Write) || event.Has(fsnotify.Remove) {
				// Now first read and see if any modification
				data, err := os.ReadFile("/etc/resolv.conf")
				if err == nil {
					if string(data) != resolvconf {
						// Now we have a different file on disk, so rewrite
						file, err := os.Create("/etc/resolv.conf")
						if err != nil {
							log.Panic("Can not write the file.")
						}
						defer file.Close()
						file.WriteString(resolvconf)
					}
				} else {
					fmt.Println("We could not read /etc/resolv.conf")
					log.Println(err)
				}

			}
		case err, ok := <-watcher.Errors:
			if !ok {
				return
			}
			log.Println("error:", err)
		}
	}
}

func Listen(port int, serveraddr string) {

	resolvconf = `# Generated by cdns
nameserver 127.0.0.1
`
	// Now we will monitor the /etc/resolv.conf

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatal(err)
	}
	defer watcher.Close()
	err = watcher.Add("/etc/resolv.conf")
	if err != nil {
		log.Fatal(err)
	}
	go watchDNSConfiguration(*watcher)
	// End of monitoring code

	// Let us read all the blocklists
	fillBlockList()

	serverurl = serveraddr
	serveMux := dns.NewServeMux()
	serveMux.HandleFunc(".", func(w dns.ResponseWriter, req *dns.Msg) {
		handleRequest(w, req)
	})
	c = cache.New(5*time.Minute, 10*time.Minute)
	question = cache.New(2*time.Second, 2*time.Minute)

	// Now for redis
	ctx = context.Background()
	// TODO: Get this from configuration file
	redisdb = redis.NewClient(&redis.Options{
		Addr:     "127.0.0.1:6379",
		Password: "", // no password set
		DB:       0,  // use default DB
	})

	server := &dns.Server{Addr: fmt.Sprintf("0.0.0.0:%d", port), Net: "udp", Handler: serveMux}
	err = server.ListenAndServe()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error while starting the server: %s\n", err)
		os.Exit(127)
	}

}

func handleRequest(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)
	if r.MsgHdr.Opcode == dns.OpcodeQuery {
		if len(r.Question) > 0 {

			// Record the query
			go recordQuery(r)

			// Check if blocklisted domain
			blocked := blocklist[r.Question[0].Name]

			if !blocked {

				// First check the cache
				res, found := c.Get(r.Question[0].String())
				if found {

					//fmt.Printf("From cache: %+v\n", r.Question)
					m.Answer = append(m.Answer, res.(dns.Msg).Answer...)
					m.Ns = append(m.Ns, res.(dns.Msg).Ns...)
					m.Extra = append(m.Extra, res.(dns.Msg).Extra...)
					w.WriteMsg(m)
					return
				}
				// Check if we already asked and waiting for an Answer
				_, found = question.Get(r.Question[0].String())
				if found {
					// Then don't ask again, fail happily
					// TODO: I don't know how bad is this idea :)
					// Set NXDOMAIN
					m.SetRcode(r, dns.RcodeNameError)
					w.WriteMsg(m)
					return
				}
			}

			// Now do a real query
			d := net.Dialer{Timeout: 500 * time.Millisecond}
			conn, err := d.Dial("udp", serverurl)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %s\n", err)
				return
			}
			defer conn.Close()
			// Now we record that we asked this question
			question.Set(r.Question[0].String(), true, cache.DefaultExpiration)
			//fmt.Printf("%+v\n", r.Question)
			dnsConn := &dns.Conn{Conn: conn}
			if err = dnsConn.WriteMsg(r); err != nil {
				// Set NXDOMAIN
				m.SetRcode(r, dns.RcodeNameError)
				w.WriteMsg(m)
				fmt.Fprintf(os.Stderr, "Error while talking to the server %s\n", err)
				return
			}
			resp, err := dnsConn.ReadMsg()
			if err == nil {
				// First delete from the question
				question.Delete(r.Question[0].String())
				// Now if the domain is not blocked then return real answer
				if !blocked {
					m.Answer = append(m.Answer, resp.Answer...)
					m.Ns = append(m.Ns, resp.Ns...)
					m.Extra = append(m.Extra, resp.Extra...)
				} else {
					// Means blocked domain, so return NXDOMAIN
					m.SetRcode(r, dns.RcodeNameError)
				}

				// We need to record the real IP address anyway
				if len(resp.Answer) > 0 {
					//fmt.Printf("TTL: %+v\n", resp.Answer[0].Header().Ttl)
					c.Set(r.Question[0].String(), *resp, cache.DefaultExpiration)
					go pushToRedis(r, resp)
				}
			}
		}
	}

	w.WriteMsg(m)
}

func pushToRedis(r *dns.Msg, answer *dns.Msg) {

	// For each IP, record the DNS name
	for _, ans := range answer.Answer {
		if ip, ok := ans.(*dns.A); ok {
			rname := fmt.Sprintf("ip:%s", ip.A.String())
			redisdb.SAdd(ctx, rname, r.Question[0].Name)
		}
		if ip, ok := ans.(*dns.AAAA); ok {
			rname := fmt.Sprintf("ip:%s", ip.AAAA.String())
			redisdb.SAdd(ctx, rname, r.Question[0].Name)
		}
	}

}

func recordQuery(r *dns.Msg) {
	name := r.Question[0].Name
	qtype := r.Question[0].Qtype
	now := time.Now()
	value := now.Unix()
	entry := fmt.Sprintf("%s:%d:%d", name, qtype, value)
	redisdb.RPush(ctx, "cdns:queries", entry)
}