-
Notifications
You must be signed in to change notification settings - Fork 0
/
iptable.go
137 lines (119 loc) · 2.63 KB
/
iptable.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
/*
* NeoDB
*
* Copyright 2018 The Radon Authors.
* Copyright 2021-2030 The NeoDB Authors.
* Code is licensed under the GPLv3.
*
*/
package proxy
import (
"regexp"
"sync"
"github.com/sealdb/neodb/config"
"github.com/sealdb/mysqlstack/xlog"
)
// IP tuple.
type IP struct {
ip string
}
// IPTable tuple.
type IPTable struct {
mu sync.RWMutex
log *xlog.Log
conf *config.ProxyConfig
iptable map[string]interface{}
}
// NewIPTable creates a new IPTable.
func NewIPTable(log *xlog.Log, conf *config.ProxyConfig) *IPTable {
ipt := &IPTable{
log: log,
conf: conf,
iptable: make(map[string]interface{}),
}
if conf.IPS != nil {
for _, ip := range conf.IPS {
if err := addToIPTable(ipt, ip); err != nil {
log.Error("add ip failed during new iptable: ip[%s], err[%+v]", ip, err)
}
}
}
return ipt
}
// isWildcardIP used to judge if ip is a regexp and return true, otherwise return false
func isWildcardIP(ip string) bool {
return regexp.QuoteMeta(ip) != ip
}
// addToIPTable is used to add an ip to iptable
func addToIPTable(ipt *IPTable, ip string) error {
if isWildcardIP(ip) {
if ip == "*" {
ip = "." + ip
}
reg, err := regexp.Compile(ip)
if err != nil {
return err
}
ipt.iptable[ip] = reg
} else {
IP := &IP{ip: ip}
ipt.iptable[ip] = IP
}
return nil
}
// Add used to add an ip to iptable.
func (ipt *IPTable) Add(ip string) error {
ipt.log.Warning("proxy.iptable.add:%s", ip)
ipt.mu.Lock()
defer ipt.mu.Unlock()
return addToIPTable(ipt, ip)
}
// Remove used to remove a ip from table.
func (ipt *IPTable) Remove(ip string) {
ipt.log.Warning("proxy.iptable.remove:%s", ip)
ipt.mu.Lock()
defer ipt.mu.Unlock()
if ip == "*" {
ip = "." + ip
}
delete(ipt.iptable, ip)
}
// Refresh used to refresh the table.
func (ipt *IPTable) Refresh() error {
ipt.log.Warning("proxy.iptable.refresh:%+v", ipt.conf.IPS)
ipt.mu.Lock()
defer ipt.mu.Unlock()
ipt.iptable = make(map[string]interface{})
if ipt.conf.IPS != nil {
for _, ip := range ipt.conf.IPS {
if err := addToIPTable(ipt, ip); err != nil {
return err
}
}
}
return nil
}
// Check used to check whether the ip is in ip table or not.
func (ipt *IPTable) Check(address string) bool {
ipt.mu.Lock()
defer ipt.mu.Unlock()
// Pass if no iptable setting.
if len(ipt.iptable) == 0 {
return true
}
// if address is in iptable[address], just return
_, ok := ipt.iptable[address]
if ok {
return true
}
// check if address is match with ip regexp
for _, ip := range ipt.iptable {
switch ip.(type) {
case *regexp.Regexp:
if ip.(*regexp.Regexp).MatchString(address) {
return true
}
}
}
return false
}