diff --git a/resolver/resolver.go b/resolver/resolver.go new file mode 100644 index 0000000..885d4fd --- /dev/null +++ b/resolver/resolver.go @@ -0,0 +1,97 @@ +// Package dns_resolver is a simple dns resolver +// based on miekg/dns +package resolver + +import ( + "errors" + "math/rand" + "net" + "os" + "strings" + "time" + + "github.com/miekg/dns" +) + +// DnsResolver represents a dns resolver +type DnsResolver struct { + Servers []string + RetryTimes int + r *rand.Rand +} + +// New initializes DnsResolver. +func New(servers []string) *DnsResolver { + for i := range servers { + servers[i] = net.JoinHostPort(servers[i], "53") + } + + return &DnsResolver{servers, len(servers) * 2, rand.New(rand.NewSource(time.Now().UnixNano()))} +} + +// NewFromResolvConf initializes DnsResolver from resolv.conf like file. +func NewFromResolvConf(path string) (*DnsResolver, error) { + if _, err := os.Stat(path); os.IsNotExist(err) { + return &DnsResolver{}, errors.New("no such file or directory: " + path) + } + config, err := dns.ClientConfigFromFile(path) + servers := []string{} + for _, ipAddress := range config.Servers { + servers = append(servers, net.JoinHostPort(ipAddress, "53")) + } + return &DnsResolver{servers, len(servers) * 2, rand.New(rand.NewSource(time.Now().UnixNano()))}, err +} + +// LookupHost returns IP addresses of provied host. +// In case of timeout retries query RetryTimes times. +func (r *DnsResolver) LookupHost(host string, dnsType uint16) ([]net.IP, error) { + return r.lookupHost(host, dnsType, r.RetryTimes) +} + +func (r *DnsResolver) lookupHost(host string, dnsType uint16, triesLeft int) ([]net.IP, error) { + m1 := new(dns.Msg) + m1.Id = dns.Id() + m1.RecursionDesired = true + m1.Question = make([]dns.Question, 1) + + switch dnsType { + case dns.TypeA: + m1.Question[0] = dns.Question{dns.Fqdn(host), dns.TypeA, dns.ClassINET} + case dns.TypeAAAA: + m1.Question[0] = dns.Question{dns.Fqdn(host), dns.TypeAAAA, dns.ClassINET} + } + + in, err := dns.Exchange(m1, r.Servers[r.r.Intn(len(r.Servers))]) + + result := []net.IP{} + + if err != nil { + if strings.HasSuffix(err.Error(), "i/o timeout") && triesLeft > 0 { + triesLeft-- + return r.lookupHost(host, dnsType, triesLeft) + } + return result, err + } + + if in != nil && in.Rcode != dns.RcodeSuccess { + return result, errors.New(dns.RcodeToString[in.Rcode]) + } + + if dnsType == dns.TypeA { + for _, record := range in.Answer { + if t, ok := record.(*dns.A); ok { + result = append(result, t.A) + } + } + } + + if dnsType == dns.TypeAAAA { + for _, record := range in.Answer { + if t, ok := record.(*dns.AAAA); ok { + result = append(result, t.AAAA) + } + } + } + + return result, err +} diff --git a/resolver/resolver_test.go b/resolver/resolver_test.go new file mode 100644 index 0000000..75d95a6 --- /dev/null +++ b/resolver/resolver_test.go @@ -0,0 +1,45 @@ +package resolver + +import ( + "fmt" + "reflect" + "testing" + + "github.com/miekg/dns" +) + +func TestNew(t *testing.T) { + servers := []string{"8.8.8.8", "8.8.4.4"} + expectedServers := []string{"8.8.8.8:53", "8.8.4.4:53"} + resolver := New(servers) + + if !reflect.DeepEqual(resolver.Servers, expectedServers) { + t.Error("resolver.Servers: ", resolver.Servers, "should be equal to", expectedServers) + } +} + +func TestLookupHost_ValidServer(t *testing.T) { + resolver := New([]string{"8.8.8.8", "8.8.4.4"}) + result, err := resolver.LookupHost("google-public-dns-a.google.com", dns.TypeA) + if err != nil { + fmt.Println(err.Error()) + t.Error("Should succeed dns lookup") + } + + if result[0].String() != "8.8.8.8" { + t.Error("google-public-dns-a.google.com should be resolved to 8.8.8.8") + } +} + +func TestLookupHostIPv6_ValidServer(t *testing.T) { + resolver := New([]string{"2001:4860:4860::8888", "2001:4860:4860::8844"}) + result, err := resolver.LookupHost("google-public-dns-a.google.com", dns.TypeAAAA) + if err != nil { + fmt.Println(err.Error()) + t.Error("Should succeed dns lookup") + } + + if result[0].String() != "2001:4860:4860::8888" { + t.Error("result shoudl be: 2001:4860:4860::8888") + } +}