mirror of
https://github.com/sot-tech/mochi.git
synced 2026-07-28 01:58:12 -07:00
store: added IPStore and memory implementation
This commit is contained in:
Leo Balduf
committed by
Jimmy Zelinskie
Jimmy Zelinskie
parent
3b54069a1b
commit
1dab3978fb
@@ -0,0 +1,159 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package memory
|
||||
|
||||
import (
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"github.com/mrd0ll4r/netmatch"
|
||||
|
||||
"github.com/chihaya/chihaya/server/store"
|
||||
)
|
||||
|
||||
func init() {
|
||||
store.RegisterIPStoreDriver("memory", &ipStoreDriver{})
|
||||
}
|
||||
|
||||
type ipStoreDriver struct{}
|
||||
|
||||
func (d *ipStoreDriver) New(cfg *store.Config) (store.IPStore, error) {
|
||||
return &ipStore{
|
||||
ips: make(map[[16]byte]struct{}),
|
||||
networks: netmatch.New(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ipStore implements store.IPStore using an in-memory map of byte arrays and
|
||||
// a trie-like structure.
|
||||
type ipStore struct {
|
||||
ips map[[16]byte]struct{}
|
||||
networks *netmatch.Trie
|
||||
sync.RWMutex
|
||||
}
|
||||
|
||||
var (
|
||||
_ store.IPStore = &ipStore{}
|
||||
v4InV6Prefix = []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff}
|
||||
)
|
||||
|
||||
// key converts an IP address to a [16]byte.
|
||||
// The byte array can then be used as a key for a map, unlike net.IP, which is a
|
||||
// []byte.
|
||||
// If an IPv4 address is specified, it will be prefixed with
|
||||
// the net.v4InV6Prefix and thus becomes a valid IPv6 address.
|
||||
func key(ip net.IP) [16]byte {
|
||||
var array [16]byte
|
||||
|
||||
if len(ip) == net.IPv4len {
|
||||
copy(array[:], v4InV6Prefix)
|
||||
copy(array[12:], ip)
|
||||
} else {
|
||||
copy(array[:], ip)
|
||||
}
|
||||
return array
|
||||
}
|
||||
|
||||
func (s *ipStore) AddNetwork(network string) error {
|
||||
s.Lock()
|
||||
defer s.Unlock()
|
||||
|
||||
key, length, err := netmatch.ParseNetwork(network)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return s.networks.Add(key, length)
|
||||
}
|
||||
|
||||
func (s *ipStore) AddIP(ip net.IP) error {
|
||||
s.Lock()
|
||||
defer s.Unlock()
|
||||
|
||||
s.ips[key(ip)] = struct{}{}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *ipStore) HasIP(ip net.IP) (bool, error) {
|
||||
s.RLock()
|
||||
defer s.RUnlock()
|
||||
key := key(ip)
|
||||
|
||||
_, ok := s.ips[key]
|
||||
if ok {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
match, err := s.networks.Match(key)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
return match, nil
|
||||
}
|
||||
|
||||
func (s *ipStore) HasAnyIP(ips []net.IP) (bool, error) {
|
||||
s.RLock()
|
||||
defer s.RUnlock()
|
||||
|
||||
for _, ip := range ips {
|
||||
key := key(ip)
|
||||
if _, ok := s.ips[key]; ok {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
match, err := s.networks.Match(key)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if match {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (s *ipStore) HasAllIPs(ips []net.IP) (bool, error) {
|
||||
s.RLock()
|
||||
defer s.RUnlock()
|
||||
|
||||
for _, ip := range ips {
|
||||
key := key(ip)
|
||||
if _, ok := s.ips[key]; !ok {
|
||||
match, err := s.networks.Match(key)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if !match {
|
||||
return false, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (s *ipStore) RemoveIP(ip net.IP) error {
|
||||
s.Lock()
|
||||
defer s.Unlock()
|
||||
|
||||
delete(s.ips, key(ip))
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *ipStore) RemoveNetwork(network string) error {
|
||||
s.Lock()
|
||||
defer s.Unlock()
|
||||
|
||||
key, length, err := netmatch.ParseNetwork(network)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return s.networks.Remove(key, length)
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package memory
|
||||
|
||||
import (
|
||||
"net"
|
||||
"testing"
|
||||
|
||||
"github.com/chihaya/chihaya/server/store"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
var (
|
||||
v6 = net.ParseIP("0c22:384e:0:0c22:384e::68")
|
||||
v4 = net.ParseIP("12.13.14.15")
|
||||
v4s = net.ParseIP("12.13.14.15").To4()
|
||||
)
|
||||
|
||||
func TestKey(t *testing.T) {
|
||||
var table = []struct {
|
||||
input net.IP
|
||||
expected [16]byte
|
||||
}{
|
||||
{v6, [16]byte{12, 34, 56, 78, 0, 0, 12, 34, 56, 78, 0, 0, 0, 0, 0, 104}},
|
||||
{v4, [16]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 255, 255, 12, 13, 14, 15}}, // IPv4 in IPv6 prefix
|
||||
{v4s, [16]byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 255, 255, 12, 13, 14, 15}}, // is equal to the one above, should produce equal output
|
||||
}
|
||||
|
||||
for _, tt := range table {
|
||||
got := key(tt.input)
|
||||
assert.Equal(t, got, tt.expected)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIPStore(t *testing.T) {
|
||||
var d = &ipStoreDriver{}
|
||||
|
||||
s, err := d.New(&store.Config{})
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, s)
|
||||
|
||||
// check default state
|
||||
found, err := s.HasIP(v4)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, found)
|
||||
|
||||
// check IPv4
|
||||
err = s.AddIP(v4)
|
||||
assert.Nil(t, err)
|
||||
|
||||
found, err = s.HasIP(v4)
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, found)
|
||||
|
||||
found, err = s.HasIP(v4s)
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, found)
|
||||
|
||||
found, err = s.HasIP(v6)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, found)
|
||||
|
||||
// check removes
|
||||
err = s.RemoveIP(v6)
|
||||
assert.Nil(t, err)
|
||||
|
||||
err = s.RemoveIP(v4s)
|
||||
assert.Nil(t, err)
|
||||
|
||||
found, err = s.HasIP(v4)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, found)
|
||||
|
||||
// check IPv6
|
||||
err = s.AddIP(v6)
|
||||
assert.Nil(t, err)
|
||||
|
||||
found, err = s.HasIP(v6)
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, found)
|
||||
|
||||
err = s.RemoveIP(v6)
|
||||
assert.Nil(t, err)
|
||||
|
||||
found, err = s.HasIP(v6)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, found)
|
||||
}
|
||||
|
||||
func TestHasAllHasAny(t *testing.T) {
|
||||
var d = &ipStoreDriver{}
|
||||
s, err := d.New(&store.Config{})
|
||||
assert.Nil(t, err)
|
||||
assert.NotNil(t, s)
|
||||
|
||||
found, err := s.HasAnyIP(nil)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, found)
|
||||
|
||||
found, err = s.HasAllIPs(nil)
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, found)
|
||||
|
||||
found, err = s.HasAllIPs([]net.IP{v4})
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, found)
|
||||
|
||||
err = s.AddIP(v4)
|
||||
assert.Nil(t, err)
|
||||
|
||||
found, err = s.HasAnyIP([]net.IP{v4, v6})
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, found)
|
||||
|
||||
found, err = s.HasAllIPs([]net.IP{v4, v6})
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, found)
|
||||
|
||||
found, err = s.HasAllIPs([]net.IP{v4})
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, found)
|
||||
|
||||
err = s.AddIP(v6)
|
||||
assert.Nil(t, err)
|
||||
|
||||
found, err = s.HasAnyIP([]net.IP{v4, v6})
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, found)
|
||||
|
||||
found, err = s.HasAllIPs([]net.IP{v4, v6})
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, found)
|
||||
}
|
||||
|
||||
func TestNetworks(t *testing.T) {
|
||||
var (
|
||||
d = &ipStoreDriver{}
|
||||
net1 = "192.168.22.255/24"
|
||||
net2 = "192.168.23.255/24"
|
||||
includedIP = net.ParseIP("192.168.22.23")
|
||||
excludedIP = net.ParseIP("192.168.23.22")
|
||||
)
|
||||
|
||||
s, err := d.New(&store.Config{})
|
||||
assert.Nil(t, err)
|
||||
|
||||
match, err := s.HasIP(includedIP)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
match, err = s.HasIP(excludedIP)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
err = s.AddNetwork("")
|
||||
assert.NotNil(t, err)
|
||||
|
||||
err = s.RemoveNetwork("")
|
||||
assert.NotNil(t, err)
|
||||
|
||||
err = s.AddNetwork(net1)
|
||||
assert.Nil(t, err)
|
||||
|
||||
match, err = s.HasIP(includedIP)
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, match)
|
||||
|
||||
match, err = s.HasIP(excludedIP)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
err = s.RemoveNetwork(net2)
|
||||
assert.NotNil(t, err)
|
||||
|
||||
err = s.RemoveNetwork(net1)
|
||||
assert.Nil(t, err)
|
||||
|
||||
match, err = s.HasIP(includedIP)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
match, err = s.HasIP(excludedIP)
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
}
|
||||
|
||||
func TestHasAllHasAnyNetworks(t *testing.T) {
|
||||
var (
|
||||
d = &ipStoreDriver{}
|
||||
net1 = "192.168.22.255/24"
|
||||
net2 = "192.168.23.255/24"
|
||||
inNet1 = net.ParseIP("192.168.22.234")
|
||||
inNet2 = net.ParseIP("192.168.23.123")
|
||||
excluded = net.ParseIP("10.154.243.22")
|
||||
)
|
||||
s, err := d.New(&store.Config{})
|
||||
assert.Nil(t, err)
|
||||
|
||||
match, err := s.HasAnyIP([]net.IP{inNet1, inNet2, excluded})
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
match, err = s.HasAllIPs([]net.IP{inNet1, inNet2, excluded})
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
err = s.AddNetwork(net1)
|
||||
assert.Nil(t, err)
|
||||
|
||||
match, err = s.HasAnyIP([]net.IP{inNet1, inNet2})
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, match)
|
||||
|
||||
match, err = s.HasAllIPs([]net.IP{inNet1, inNet2})
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
err = s.AddNetwork(net2)
|
||||
assert.Nil(t, err)
|
||||
|
||||
match, err = s.HasAnyIP([]net.IP{inNet1, inNet2, excluded})
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, match)
|
||||
|
||||
match, err = s.HasAllIPs([]net.IP{inNet1, inNet2})
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, match)
|
||||
|
||||
match, err = s.HasAllIPs([]net.IP{inNet1, inNet2, excluded})
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
err = s.RemoveNetwork(net1)
|
||||
assert.Nil(t, err)
|
||||
|
||||
match, err = s.HasAnyIP([]net.IP{inNet1, inNet2})
|
||||
assert.Nil(t, err)
|
||||
assert.True(t, match)
|
||||
|
||||
match, err = s.HasAllIPs([]net.IP{inNet1, inNet2})
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
err = s.RemoveNetwork(net2)
|
||||
assert.Nil(t, err)
|
||||
|
||||
match, err = s.HasAnyIP([]net.IP{inNet1, inNet2})
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
|
||||
match, err = s.HasAllIPs([]net.IP{inNet1, inNet2})
|
||||
assert.Nil(t, err)
|
||||
assert.False(t, match)
|
||||
}
|
||||
Reference in New Issue
Block a user