github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/trie/cidr/cidr.go (about)

     1  package cidr
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"net/netip"
     7  )
     8  
     9  // Cidr cidr matcher
    10  type Cidr[T any] struct {
    11  	v4CidrTrie Trie[T]
    12  	v6CidrTrie Trie[T]
    13  }
    14  
    15  // InsetOneCIDR Insert one CIDR to cidr matcher
    16  func (c *Cidr[T]) Insert(cidr string, mark T) error {
    17  	ipNet, err := netip.ParsePrefix(cidr)
    18  	if err != nil {
    19  		return fmt.Errorf("parse cidr [%s] failed: %w", cidr, err)
    20  	}
    21  	c.InsertCIDR(ipNet, mark)
    22  	return nil
    23  }
    24  
    25  func (c *Cidr[T]) RemoveCIDR(ipNet netip.Prefix) {
    26  	//TODO
    27  }
    28  func (c *Cidr[T]) RemoveIP(ipNet netip.Addr, maskSIze int) {
    29  	//TODO
    30  }
    31  
    32  func (c *Cidr[T]) InsertCIDR(ipNet netip.Prefix, mark T) {
    33  	if ipNet.Addr().Is4() {
    34  		c.v4CidrTrie.Insert(ipNet.Addr().AsSlice(), ipNet.Bits(), mark)
    35  	} else {
    36  		c.v6CidrTrie.Insert(ipNet.Addr().AsSlice(), ipNet.Bits(), mark)
    37  	}
    38  }
    39  
    40  func (c *Cidr[T]) InsertIP(ip netip.Addr, maskSize int, mark T) {
    41  	if ip.Is4() {
    42  		c.v4CidrTrie.Insert(ip.AsSlice(), maskSize, mark)
    43  	} else {
    44  		c.v6CidrTrie.Insert(ip.AsSlice(), maskSize, mark)
    45  	}
    46  }
    47  
    48  // MatchWithTrie match ip with trie
    49  func (c *Cidr[T]) Search(ip string) (mark T, ok bool) {
    50  	iP := net.ParseIP(ip)
    51  	if iP == nil {
    52  		return mark, false
    53  	}
    54  	return c.SearchIP(iP)
    55  }
    56  
    57  func (c *Cidr[T]) SearchIP(ip net.IP) (mark T, ok bool) {
    58  	if x := ip.To4(); x != nil {
    59  		return c.v4CidrTrie.Search(x)
    60  	} else {
    61  		return c.v6CidrTrie.Search(ip.To16())
    62  	}
    63  }
    64  
    65  func NewCidrMapper[T any]() *Cidr[T] {
    66  	cidrMapper := new(Cidr[T])
    67  	cidrMapper.v4CidrTrie = NewTrieTree[T]()
    68  	cidrMapper.v6CidrTrie = NewTrieTree[T]()
    69  	return cidrMapper
    70  }