diff --git a/config/example.yaml b/config/example.yaml index 88815aa..e5fa5c1 100644 --- a/config/example.yaml +++ b/config/example.yaml @@ -21,6 +21,7 @@ servers: readTimeout: 10s writeTimeout: 10s clientStore: memory + ipStore: memory peerStore: memory peerStoreConfig: gcAfter: 30m diff --git a/server/store/ip_store.go b/server/store/ip_store.go new file mode 100644 index 0000000..70a4fad --- /dev/null +++ b/server/store/ip_store.go @@ -0,0 +1,79 @@ +// 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 store + +import ( + "fmt" + "net" +) + +var ipStoreDrivers = make(map[string]IPStoreDriver) + +// IPStore represents an interface for manipulating IPs and IP ranges. +type IPStore interface { + // AddIP adds a single IP address to the IPStore. + AddIP(ip net.IP) error + + // AddNetwork adds a range of IP addresses, denoted by a network in CIDR + // notation, to the IPStore. + AddNetwork(network string) error + + // HasIP returns whether the given IP address is contained in the IPStore + // or belong to any of the stored networks. + HasIP(ip net.IP) (bool, error) + + // HasAnyIP returns whether any of the given IP addresses are contained in + // the IPStore or belong to any of the stored networks. + HasAnyIP(ips []net.IP) (bool, error) + + // HassAllIPs returns whether all of the given IP addresses are contained in + // the IPStore or belong to any of the stored networks. + HasAllIPs(ips []net.IP) (bool, error) + + // RemoveIP removes a single IP address from the IPStore. + // + // This wil not remove the given address from any networks it belongs to + // that are stored in the IPStore. + RemoveIP(ip net.IP) error + + // RemoveNetwork removes a range of IP addresses that was previously added + // through AddNetwork. + // + // The given network must not, as a string, match the previously added + // network, but rather denote the same network, e.g. if the network + // 192.168.22.255/24 was added, removing the network 192.168.22.123/24 + // will succeed. + RemoveNetwork(network string) error +} + +// IPStoreDriver represents an interface for creating a handle to the +// storage of IPs. +type IPStoreDriver interface { + New(*Config) (IPStore, error) +} + +// RegisterIPStoreDriver makes a driver available by the provided name. +// +// If this function is called twice with the same name or if the driver is nil, +// it panics. +func RegisterIPStoreDriver(name string, driver IPStoreDriver) { + if driver == nil { + panic("store: could not register nil ClientStoreDriver") + } + if _, dup := ipStoreDrivers[name]; dup { + panic("store: could not register duplicate ClientStoreDriver: " + name) + } + ipStoreDrivers[name] = driver +} + +// OpenIPStore returns an IPStore specified by a configuration. +func OpenIPStore(cfg *Config) (IPStore, error) { + driver, ok := ipStoreDrivers[cfg.IPStore] + if !ok { + return nil, fmt.Errorf("store: unknown driver %q (forgotten import?)", cfg.IPStore) + } + + return driver.New(cfg) +} diff --git a/server/store/memory/ip_store.go b/server/store/memory/ip_store.go new file mode 100644 index 0000000..be91a12 --- /dev/null +++ b/server/store/memory/ip_store.go @@ -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) +} diff --git a/server/store/memory/ip_store_test.go b/server/store/memory/ip_store_test.go new file mode 100644 index 0000000..73175e9 --- /dev/null +++ b/server/store/memory/ip_store_test.go @@ -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) +} diff --git a/server/store/store.go b/server/store/store.go index c50c4d7..9ed2d0b 100644 --- a/server/store/store.go +++ b/server/store/store.go @@ -48,6 +48,8 @@ type Config struct { ClientStoreConfig interface{} `yaml:"clienStoreConfig"` PeerStore string `yaml:"peerStore"` PeerStoreConfig interface{} `yaml:"peerStoreConfig"` + IPStore string `yaml:"ipStore"` + IPStoreConfig interface{} `yaml:"ipStoreConfig"` } func newConfig(srvcfg interface{}) (*Config, error) { @@ -84,6 +86,7 @@ type Store struct { PeerStore ClientStore + IPStore } func (s *Store) Start() {