ipcidr_set.go 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. package cidr
  2. import (
  3. "fmt"
  4. "net/netip"
  5. "unsafe"
  6. "go4.org/netipx"
  7. )
  8. type IpCidrSet struct {
  9. // must same with netipx.IPSet
  10. rr []netipx.IPRange
  11. }
  12. func NewIpCidrSet() *IpCidrSet {
  13. return &IpCidrSet{}
  14. }
  15. func (set *IpCidrSet) AddIpCidrForString(ipCidr string) error {
  16. prefix, err := netip.ParsePrefix(ipCidr)
  17. if err != nil {
  18. return err
  19. }
  20. return set.AddIpCidr(prefix)
  21. }
  22. func (set *IpCidrSet) AddIpCidr(ipCidr netip.Prefix) (err error) {
  23. if r := netipx.RangeOfPrefix(ipCidr); r.IsValid() {
  24. set.rr = append(set.rr, r)
  25. } else {
  26. err = fmt.Errorf("not valid ipcidr range: %s", ipCidr)
  27. }
  28. return
  29. }
  30. func (set *IpCidrSet) IsContainForString(ipString string) bool {
  31. ip, err := netip.ParseAddr(ipString)
  32. if err != nil {
  33. return false
  34. }
  35. return set.IsContain(ip)
  36. }
  37. func (set *IpCidrSet) IsContain(ip netip.Addr) bool {
  38. return set.ToIPSet().Contains(ip.WithZone(""))
  39. }
  40. func (set *IpCidrSet) Merge() error {
  41. var b netipx.IPSetBuilder
  42. b.AddSet(set.ToIPSet())
  43. i, err := b.IPSet()
  44. if err != nil {
  45. return err
  46. }
  47. set.fromIPSet(i)
  48. return nil
  49. }
  50. func (set *IpCidrSet) Foreach(f func(prefix netip.Prefix) bool) {
  51. for _, r := range set.rr {
  52. for _, prefix := range r.Prefixes() {
  53. if !f(prefix) {
  54. return
  55. }
  56. }
  57. }
  58. }
  59. // ToIPSet not safe convert to *netipx.IPSet
  60. // be careful, must be used after Merge
  61. func (set *IpCidrSet) ToIPSet() *netipx.IPSet {
  62. return (*netipx.IPSet)(unsafe.Pointer(set))
  63. }
  64. func (set *IpCidrSet) fromIPSet(i *netipx.IPSet) {
  65. *set = *(*IpCidrSet)(unsafe.Pointer(i))
  66. }