forked from grafana/k6
-
Notifications
You must be signed in to change notification settings - Fork 0
/
resolver.go
63 lines (56 loc) · 1.43 KB
/
resolver.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
package mockresolver
import (
"fmt"
"net"
"sync"
)
// MockResolver implements netext.Resolver, and allows changing the host
// mapping at runtime.
type MockResolver struct {
m sync.RWMutex
hosts map[string][]net.IP
ResolveHook func(host string, result []net.IP)
}
// New returns a new MockResolver.
func New(hosts map[string][]net.IP) *MockResolver {
if hosts == nil {
hosts = make(map[string][]net.IP)
}
return &MockResolver{hosts: hosts}
}
// LookupIP returns the first IP mapped for host.
func (r *MockResolver) LookupIP(host string) (net.IP, error) {
if ips, err := r.LookupIPAll(host); err != nil {
return nil, err
} else if len(ips) > 0 {
return ips[0], nil
}
return nil, nil
}
// LookupIPAll returns all IPs mapped for host. It mimics the net.LookupIP
// signature so that it can be used to mock netext.LookupIP in tests.
func (r *MockResolver) LookupIPAll(host string) (ips []net.IP, err error) {
defer func() {
if r.ResolveHook != nil {
r.ResolveHook(host, ips)
}
}()
r.m.RLock()
defer r.m.RUnlock()
if ips, ok := r.hosts[host]; ok {
return ips, nil
}
return nil, fmt.Errorf("lookup %s: no such host", host)
}
// Set the host to resolve to ip.
func (r *MockResolver) Set(host, ip string) {
r.m.Lock()
defer r.m.Unlock()
r.hosts[host] = []net.IP{net.ParseIP(ip)}
}
// Unset removes the host.
func (r *MockResolver) Unset(host string) {
r.m.Lock()
defer r.m.Unlock()
delete(r.hosts, host)
}