-
Notifications
You must be signed in to change notification settings - Fork 4
/
pool.go
160 lines (142 loc) · 3.02 KB
/
pool.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
151
152
153
154
155
156
157
158
159
160
package fakeip
import (
"net"
"sync"
"time"
"github.com/nickjfree/goose/pkg/routing/rule"
"github.com/nickjfree/goose/pkg/utils"
)
// dns record A
type dnsRecord struct {
IP net.IP
Exp time.Time
}
// fake ip manager
type FakeIPManager struct {
// network
network net.IPNet
// fake ip pool
pool *utils.IPPool
// fake to real ip mapping
f2r *utils.IPMapping
// real to fake ip mapping
r2f *utils.IPMapping
// fakeip rule
rule *rule.Rule
// fale name server cache
nameServer map[string][]dnsRecord
// lock
mu sync.Mutex
}
func NewFakeIPManager(network, script, db string) *FakeIPManager {
_, ipNet, err := net.ParseCIDR(network)
if err != nil {
logger.Fatal(err)
}
pool := utils.NewIPPool(*ipNet)
m := &FakeIPManager{
network: *ipNet,
pool: pool,
f2r: utils.NewIPMapping(func(ip net.IP) error {
pool.Free(ip)
return nil
}),
r2f: utils.NewIPMapping(nil),
nameServer: make(map[string][]dnsRecord),
}
if script != "" && db != "" {
m.rule = rule.New(script, db)
if err := m.rule.Run(); err != nil {
logger.Fatal(err)
}
}
return m
}
// alloc fake ip
func (manager *FakeIPManager) alloc(domain string, real net.IP) (net.IP, error) {
manager.mu.Lock()
defer manager.mu.Unlock()
var fake net.IP
var err error
// find fakeip from mapping
if f := manager.r2f.Get(real); f != nil {
fake = *f
} else {
// alloc new fake ip
if fake, err = manager.pool.Alloc(); err != nil {
return nil, err
}
}
// update mapping
manager.f2r.Put(fake, real)
manager.r2f.Put(real, fake)
return fake, nil
}
// get real ip by fake ip
func (manager *FakeIPManager) toReal(fake net.IP) *net.IP {
manager.mu.Lock()
defer manager.mu.Unlock()
return manager.f2r.Get(fake)
}
// get fake ip by real ip
func (manager *FakeIPManager) toFake(real net.IP) *net.IP {
manager.mu.Lock()
defer manager.mu.Unlock()
return manager.r2f.Get(real)
}
// dns traffice routing
func (manager *FakeIPManager) DNSRoutings() []net.IPNet {
return []net.IPNet{
manager.network,
{
IP: net.IPv4(8, 8, 8, 8),
Mask: net.IPv4Mask(255, 255, 255, 255),
},
}
}
// register custome dns record to fakeip manager
func (manager *FakeIPManager) SetNameRecord(name string, ip net.IP) {
manager.mu.Lock()
defer manager.mu.Unlock()
// expire time for the record
exp := time.Now().Add(3 * time.Minute)
answers, ok := manager.nameServer[name]
if ok {
found := false
for i := range answers {
if answers[i].IP.Equal(ip) {
answers[i].Exp = exp
found = true
break
}
}
if !found {
answers = append(answers, dnsRecord{
IP: ip,
Exp: exp,
})
}
} else {
answers = []dnsRecord{
{
IP: ip,
Exp: exp,
},
}
}
manager.nameServer[name] = answers
}
func (manager *FakeIPManager) GetNameRecord(name string) []net.IP {
current := time.Now()
manager.mu.Lock()
defer manager.mu.Unlock()
result := []net.IP{}
if answers, ok := manager.nameServer[name]; ok {
for _, ans := range answers {
if current.Sub(ans.Exp) < 0 {
result = append(result, ans.IP)
}
}
}
return result
}