Checkpoint

dataset/network.go

@@ -0,0 +1,52 @@
+package dataset
+
+import (
+	"net"
+
+	"github.com/yl2chen/cidranger"
+)
+
+type NetworkTree struct {
+	ranger cidranger.Ranger
+}
+
+func MustNetworkTree(networks ...string) *NetworkTree {
+	tree, err := NewNetworkTree(networks...)
+	if err != nil {
+		panic(err)
+	}
+	return tree
+}
+
+func NewNetworkTree(networks ...string) (*NetworkTree, error) {
+	tree := &NetworkTree{
+		ranger: cidranger.NewPCTrieRanger(),
+	}
+	for _, cidr := range networks {
+		if err := tree.AddCIDR(cidr); err != nil {
+			return nil, err
+		}
+	}
+	return tree, nil
+}
+
+func (tree *NetworkTree) Add(ipnet *net.IPNet) {
+	if ipnet == nil {
+		return
+	}
+	tree.ranger.Insert(cidranger.NewBasicRangerEntry(*ipnet))
+}
+
+func (tree *NetworkTree) AddCIDR(cidr string) error {
+	_, ipnet, err := net.ParseCIDR(cidr)
+	if err != nil {
+		return err
+	}
+	tree.ranger.Insert(cidranger.NewBasicRangerEntry(*ipnet))
+	return nil
+}
+
+func (tree *NetworkTree) Contains(ip net.IP) bool {
+	contains, _ := tree.ranger.Contains(ip)
+	return contains
+}