main.go

package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/chekt/dns"
)

const resolvConfigFilePath = "/etc/resolv.conf"

type serverInfo struct {
	ServerStringForConfig string // e.g. nameserver 192.168.0.1
	Server                string // e.g. 192.168.0.1:53
	Latency               time.Duration
	Error                 error
}
type serverInfoList []serverInfo

func getNameServers() (serverInfoList serverInfoList, err error) {
	config, _ := dns.ClientConfigFromFile(resolvConfigFilePath)
	port := "53"
	if config.Port != "" {
		port = config.Port
	}
	if len(config.Servers) == 0 {
		errMessage := fmt.Sprintf("No nameservers in %s", resolvConfigFilePath)
		err = errors.New(errMessage)
	} else {
		servers := config.Servers
		for _, v := range servers {
			var serverInfo serverInfo
			server := v + ":" + port
			serverInfo.Server = server
			serverInfo.ServerStringForConfig = "nameserver " + v
			serverInfoList = append(serverInfoList, serverInfo)
		}
	}
	return serverInfoList, err
}

func checkNameServerSpeed(server string, target string) (duration time.Duration, records []string, err error) {
	c := dns.Client{}
	m := dns.Msg{}
	m.SetQuestion(target+".", dns.TypeA)
	res, duration, err := c.Exchange(&m, server)
	if err != nil {
		return
	}
	for _, answer := range res.Answer {
		records = append(records, strings.Replace(answer.String(), "\t", " ", -1))
	}
	totalRecord := len(records)
	if totalRecord == 0 {
		errMessage := fmt.Sprintf("No IP address for %s", target)
		err = errors.New(errMessage)
		return
	}
	return
}

func updateResolvconf(serverInfoList serverInfoList) (err error) {
	context := "# Generated by dnsResolverOptimizer\n"
	contextServerWithErr := ""
	for _, serverInfo := range serverInfoList {
		if serverInfo.Error != nil {
			contextServerWithErr += serverInfo.ServerStringForConfig + "\n"
			continue
		}
		context += serverInfo.ServerStringForConfig + "\n"
	}
	context = context + contextServerWithErr
	textToWrite := []byte(context)
	err = ioutil.WriteFile(resolvConfigFilePath, textToWrite, 0644)
	return err
}

func main() {

	// 1. Get target hostname
	target := "google.com"
	// server := "203.248.252.2"
	// server = "164.124.250.84"

	// 2. UPDATE argument as target
	if len(os.Args) > 1 {
		target = os.Args[1]
	}
	log.Println(target)

	// 3. Get nameserver from resolv.conf
	serverInfoList, err := getNameServers()
	if err != nil {
		log.Fatal(err)
	}

	// 4. Check nameserver latency
	for i, serverInfo := range serverInfoList {
		server := serverInfo.Server
		duration, records, err := checkNameServerSpeed(server, target)
		if err != nil {
			message := fmt.Sprintf("server: %v, records: %+v, speed: %v, error: %v", server, records, duration, err)
			log.Println(message)
		} else {
			message := fmt.Sprintf("server: %v, records: %+v, speed: %v, error: %v", server, records, duration, err)
			log.Println(message)
		}
		serverInfoList[i].Error = err
		serverInfoList[i].Latency = duration
	}

	// 5. Sort nameserver by latency
	sort.SliceStable(serverInfoList, func(i, j int) bool {
		return serverInfoList[i].Latency < serverInfoList[j].Latency
	})

	// 6. Update resolv.conf
	err = updateResolvconf(serverInfoList)
	if err != nil {
		log.Fatal(err)
	}

}