/
resolve.go
150 lines (127 loc) · 3.09 KB
/
resolve.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
package resolve
import (
"sync"
"github.com/miekg/dns"
"github.com/rs/xid"
)
const (
maxResolveRetries = 5
maxWildcardChecks = 3
)
// ResolutionPool is a pool of resolvers created for resolving subdomains
// for a given host.
type ResolutionPool struct {
*Resolver
Tasks chan string
Results chan Result
wg *sync.WaitGroup
removeWildcard bool
wildcardIPs map[string]struct{}
}
// Result contains the result for a host resolution
type Result struct {
Type ResultType
Host string
IP string
Error error
}
// ResultType is the type of result found
type ResultType int
// Types of data result can return
const (
Subdomain ResultType = iota
Error
)
// NewResolutionPool creates a pool of resolvers for resolving subdomains of a given domain
func (r *Resolver) NewResolutionPool(workers int, removeWildcard bool) *ResolutionPool {
resolutionPool := &ResolutionPool{
Resolver: r,
Tasks: make(chan string),
Results: make(chan Result),
wg: &sync.WaitGroup{},
removeWildcard: removeWildcard,
wildcardIPs: make(map[string]struct{}),
}
go func() {
for i := 0; i < workers; i++ {
resolutionPool.wg.Add(1)
go resolutionPool.resolveWorker()
}
resolutionPool.wg.Wait()
close(resolutionPool.Results)
}()
return resolutionPool
}
// InitWildcards inits the wildcard ips array
func (r *ResolutionPool) InitWildcards(domain string) error {
for i := 0; i < maxWildcardChecks; i++ {
uid := xid.New().String()
hosts, err := r.getARecords(uid + "." + domain)
if err != nil {
return err
}
// Append all wildcard ips found for domains
for _, host := range hosts {
r.wildcardIPs[host] = struct{}{}
}
}
return nil
}
func (r *ResolutionPool) resolveWorker() {
for task := range r.Tasks {
if !r.removeWildcard {
r.Results <- Result{Type: Subdomain, Host: task, IP: ""}
continue
}
hosts, err := r.getARecords(task)
if err != nil {
r.Results <- Result{Type: Error, Error: err}
continue
}
if len(hosts) == 0 {
continue
}
for _, host := range hosts {
// Ignore the host if it exists in wildcard ips map
if _, ok := r.wildcardIPs[host]; ok {
continue
}
}
r.Results <- Result{Type: Subdomain, Host: task, IP: hosts[0]}
}
r.wg.Done()
}
// getARecords gets all the A records for a given host
func (r *ResolutionPool) getARecords(host string) ([]string, error) {
var iteration int
m := new(dns.Msg)
m.Id = dns.Id()
m.RecursionDesired = true
m.Question = make([]dns.Question, 1)
m.Question[0] = dns.Question{
Name: dns.Fqdn(host),
Qtype: dns.TypeA,
Qclass: dns.ClassINET,
}
exchange:
iteration++
in, err := dns.Exchange(m, r.resolvers[r.rand.Intn(len(r.resolvers))]+":53")
if err != nil {
// Retry in case of I/O error
if iteration <= maxResolveRetries {
goto exchange
}
return nil, err
}
// Ignore the error in case we have bad result
if in != nil && in.Rcode != dns.RcodeSuccess {
return nil, nil
}
var hosts []string
for _, record := range in.Answer {
if t, ok := record.(*dns.A); ok {
hosts = append(hosts, t.A.String())
}
}
return hosts, nil
}