github.com/xmplusdev/xmcore@v1.8.11-0.20240412132628-5518b55526af/app/router/condition_geoip.go (about)

     1  package router
     2  
     3  import (
     4  	"net/netip"
     5  	"strconv"
     6  
     7  	"github.com/xmplusdev/xmcore/common/net"
     8  	"go4.org/netipx"
     9  )
    10  
    11  type GeoIPMatcher struct {
    12  	countryCode  string
    13  	reverseMatch bool
    14  	ip4          *netipx.IPSet
    15  	ip6          *netipx.IPSet
    16  }
    17  
    18  func (m *GeoIPMatcher) Init(cidrs []*CIDR) error {
    19  	var builder4, builder6 netipx.IPSetBuilder
    20  
    21  	for _, cidr := range cidrs {
    22  		ip := net.IP(cidr.GetIp())
    23  		ipPrefixString := ip.String() + "/" + strconv.Itoa(int(cidr.GetPrefix()))
    24  		ipPrefix, err := netip.ParsePrefix(ipPrefixString)
    25  		if err != nil {
    26  			return err
    27  		}
    28  
    29  		switch len(ip) {
    30  		case net.IPv4len:
    31  			builder4.AddPrefix(ipPrefix)
    32  		case net.IPv6len:
    33  			builder6.AddPrefix(ipPrefix)
    34  		}
    35  	}
    36  
    37  	if ip4, err := builder4.IPSet(); err != nil {
    38  		return err
    39  	} else {
    40  		m.ip4 = ip4
    41  	}
    42  
    43  	if ip6, err := builder6.IPSet(); err != nil {
    44  		return err
    45  	} else {
    46  		m.ip6 = ip6
    47  	}
    48  
    49  	return nil
    50  }
    51  
    52  func (m *GeoIPMatcher) SetReverseMatch(isReverseMatch bool) {
    53  	m.reverseMatch = isReverseMatch
    54  }
    55  
    56  func (m *GeoIPMatcher) match4(ip net.IP) bool {
    57  	nip, ok := netipx.FromStdIP(ip)
    58  	if !ok {
    59  		return false
    60  	}
    61  
    62  	return m.ip4.Contains(nip)
    63  }
    64  
    65  func (m *GeoIPMatcher) match6(ip net.IP) bool {
    66  	nip, ok := netipx.FromStdIP(ip)
    67  	if !ok {
    68  		return false
    69  	}
    70  
    71  	return m.ip6.Contains(nip)
    72  }
    73  
    74  // Match returns true if the given ip is included by the GeoIP.
    75  func (m *GeoIPMatcher) Match(ip net.IP) bool {
    76  	isMatched := false
    77  	switch len(ip) {
    78  	case net.IPv4len:
    79  		isMatched = m.match4(ip)
    80  	case net.IPv6len:
    81  		isMatched = m.match6(ip)
    82  	}
    83  	if m.reverseMatch {
    84  		return !isMatched
    85  	}
    86  	return isMatched
    87  }
    88  
    89  // GeoIPMatcherContainer is a container for GeoIPMatchers. It keeps unique copies of GeoIPMatcher by country code.
    90  type GeoIPMatcherContainer struct {
    91  	matchers []*GeoIPMatcher
    92  }
    93  
    94  // Add adds a new GeoIP set into the container.
    95  // If the country code of GeoIP is not empty, GeoIPMatcherContainer will try to find an existing one, instead of adding a new one.
    96  func (c *GeoIPMatcherContainer) Add(geoip *GeoIP) (*GeoIPMatcher, error) {
    97  	if len(geoip.CountryCode) > 0 {
    98  		for _, m := range c.matchers {
    99  			if m.countryCode == geoip.CountryCode && m.reverseMatch == geoip.ReverseMatch {
   100  				return m, nil
   101  			}
   102  		}
   103  	}
   104  
   105  	m := &GeoIPMatcher{
   106  		countryCode:  geoip.CountryCode,
   107  		reverseMatch: geoip.ReverseMatch,
   108  	}
   109  	if err := m.Init(geoip.Cidr); err != nil {
   110  		return nil, err
   111  	}
   112  	if len(geoip.CountryCode) > 0 {
   113  		c.matchers = append(c.matchers, m)
   114  	}
   115  	return m, nil
   116  }
   117  
   118  var globalGeoIPContainer GeoIPMatcherContainer