-
Notifications
You must be signed in to change notification settings - Fork 1.2k
/
ipblock.go
201 lines (182 loc) · 5.72 KB
/
ipblock.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
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
/*
*
* k6 - a next-generation load testing tool
* Copyright (C) 2020 Load Impact
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as
* published by the Free Software Foundation, either version 3 of the
* License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*
*/
package types
import (
"errors"
"fmt"
"math/big"
"net"
"strings"
)
// ipBlock represents a continuous segment of IP addresses
type ipBlock struct {
firstIP, count *big.Int
ipv6 bool
}
// ipPoolBlock is almost the same as ipBlock but is put in a IPPool and knows it's start id indest
// of it's size/count
type ipPoolBlock struct {
firstIP, startIndex *big.Int
}
// IPPool represent a slice of IPBlocks
type IPPool struct {
list []ipPoolBlock
count *big.Int
}
func getIPBlock(s string) (*ipBlock, error) {
switch {
case strings.Contains(s, "-"):
return ipBlockFromRange(s)
case strings.Contains(s, "/"):
return ipBlockFromCIDR(s)
default:
if net.ParseIP(s) == nil {
return nil, fmt.Errorf("%s is not a valid IP, IP range or CIDR", s)
}
return ipBlockFromRange(s + "-" + s)
}
}
func ipBlockFromRange(s string) (*ipBlock, error) {
ss := strings.SplitN(s, "-", 2)
ip0Str, ip1Str := strings.TrimSpace(ss[0]), strings.TrimSpace(ss[1])
ip0, ip1 := net.ParseIP(ip0Str), net.ParseIP(ip1Str)
if ip0 == nil || ip1 == nil {
return nil, errors.New("wrong IP range format: " + s)
}
if (ip0.To4() == nil) != (ip1.To4() == nil) { // XOR
return nil, errors.New("mixed IP range format: " + s)
}
block := ipBlockFromTwoIPs(ip0, ip1)
if block.count.Sign() <= 0 {
return nil, errors.New("negative IP range: " + s)
}
return block, nil
}
func ipBlockFromTwoIPs(ip0, ip1 net.IP) *ipBlock {
// This code doesn't do any checks on the validity of the arguments, that should be
// done before and/or after it is called
var block ipBlock
block.firstIP = new(big.Int)
block.count = new(big.Int)
block.ipv6 = ip0.To4() == nil
if block.ipv6 {
block.firstIP.SetBytes(ip0.To16())
block.count.SetBytes(ip1.To16())
} else {
block.firstIP.SetBytes(ip0.To4())
block.count.SetBytes(ip1.To4())
}
block.count.Sub(block.count, block.firstIP)
block.count.Add(block.count, big.NewInt(1))
return &block
}
func ipBlockFromCIDR(s string) (*ipBlock, error) {
_, pnet, err := net.ParseCIDR(s)
if err != nil {
return nil, fmt.Errorf("parseCIDR() failed parsing %s: %w", s, err)
}
ip0 := pnet.IP
// TODO: this is just to copy it, it will probably be better to copy the bytes ...
ip1 := net.ParseIP(ip0.String())
if ip1.To4() == nil {
ip1 = ip1.To16()
} else {
ip1 = ip1.To4()
}
for i := range ip1 {
ip1[i] |= (255 ^ pnet.Mask[i])
}
block := ipBlockFromTwoIPs(ip0, ip1)
// in the case of ipv4 if the network is bigger than 31 the first and last IP are reserved so we
// need to reduce the addresses by 2 and increment the first ip
if !block.ipv6 && big.NewInt(2).Cmp(block.count) < 0 {
block.count.Sub(block.count, big.NewInt(2))
block.firstIP.Add(block.firstIP, big.NewInt(1))
}
return block, nil
}
func (b ipPoolBlock) getIP(index *big.Int) net.IP {
// TODO implement walking ipv6 networks first
// that will probably require more math ... including knowing which is the next network and ...
// thinking about it - it looks like it's going to be kind of hard or badly defined
i := new(big.Int)
i.Add(b.firstIP, index)
// TODO use big.Int.FillBytes when golang 1.14 is no longer supported
return net.IP(i.Bytes())
}
// NewIPPool returns an IPPool slice from the provided string representation that should be comma
// separated list of IPs, IP ranges(ip1-ip2) and CIDRs
func NewIPPool(ranges string) (*IPPool, error) {
ss := strings.Split(strings.TrimSpace(ranges), ",")
pool := &IPPool{}
pool.list = make([]ipPoolBlock, len(ss))
pool.count = new(big.Int)
for i, bs := range ss {
r, err := getIPBlock(bs)
if err != nil {
return nil, err
}
pool.list[i] = ipPoolBlock{
firstIP: r.firstIP,
startIndex: new(big.Int).Set(pool.count), // this is how many there are until now
}
pool.count.Add(pool.count, r.count)
}
// The list gets reversed here as later it is searched based on when the index we are looking is
// bigger than startIndex but it will be true always for the first block which is with
// startIndex 0. This can also be fixed by iterating in reverse but this seems better
for i := 0; i < len(pool.list)/2; i++ {
pool.list[i], pool.list[len(pool.list)/2-i] = pool.list[len(pool.list)/2-i], pool.list[i]
}
return pool, nil
}
// GetIP return an IP from a pool of IPBlock slice
func (pool *IPPool) GetIP(index uint64) net.IP {
return pool.GetIPBig(new(big.Int).SetUint64(index))
}
// GetIPBig returns an IP from the pool with the provided index that is big.Int
func (pool *IPPool) GetIPBig(index *big.Int) net.IP {
index = new(big.Int).Rem(index, pool.count)
for _, b := range pool.list {
if index.Cmp(b.startIndex) >= 0 {
return b.getIP(index.Sub(index, b.startIndex))
}
}
return nil
}
// NullIPPool is a nullable IPPool
type NullIPPool struct {
Pool *IPPool
Valid bool
}
// UnmarshalText converts text data to a valid NullIPPool
func (n *NullIPPool) UnmarshalText(data []byte) error {
if len(data) == 0 {
*n = NullIPPool{}
return nil
}
var err error
n.Pool, err = NewIPPool(string(data))
if err != nil {
return err
}
n.Valid = true
return nil
}