store: added IPStore and memory implementation

config/example.yaml

@@ -21,6 +21,7 @@ servers:
       readTimeout: 10s
       writeTimeout: 10s
       clientStore: memory
+      ipStore: memory
       peerStore: memory
       peerStoreConfig:
         gcAfter: 30m

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)
+}

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)
+}

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)
+}

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() {