chore: cleanup IPSet code

This commit is contained in:
wwqgtxx 2024-01-13 11:44:02 +08:00
parent d2d8c0115c
commit e860497c0c
3 changed files with 20 additions and 190 deletions

View File

@ -1,12 +1,11 @@
package router
import (
"encoding/binary"
"fmt"
"net"
"sort"
"net/netip"
"strings"
"github.com/metacubex/mihomo/component/cidr"
"github.com/metacubex/mihomo/component/geodata/strmatcher"
"github.com/metacubex/mihomo/component/trie"
)
@ -121,115 +120,24 @@ func (m *v2rayDomainMatcher) ApplyDomain(domain string) bool {
return isMatched
}
// CIDRList is an alias of []*CIDR to provide sort.Interface.
type CIDRList []*CIDR
// Len implements sort.Interface.
func (l *CIDRList) Len() int {
return len(*l)
}
// Less implements sort.Interface.
func (l *CIDRList) Less(i int, j int) bool {
ci := (*l)[i]
cj := (*l)[j]
if len(ci.Ip) < len(cj.Ip) {
return true
}
if len(ci.Ip) > len(cj.Ip) {
return false
}
for k := 0; k < len(ci.Ip); k++ {
if ci.Ip[k] < cj.Ip[k] {
return true
}
if ci.Ip[k] > cj.Ip[k] {
return false
}
}
return ci.Prefix < cj.Prefix
}
// Swap implements sort.Interface.
func (l *CIDRList) Swap(i int, j int) {
(*l)[i], (*l)[j] = (*l)[j], (*l)[i]
}
type ipv6 struct {
a uint64
b uint64
}
type GeoIPMatcher struct {
countryCode string
reverseMatch bool
ip4 []uint32
prefix4 []uint8
ip6 []ipv6
prefix6 []uint8
}
func normalize4(ip uint32, prefix uint8) uint32 {
return (ip >> (32 - prefix)) << (32 - prefix)
}
func normalize6(ip ipv6, prefix uint8) ipv6 {
if prefix <= 64 {
ip.a = (ip.a >> (64 - prefix)) << (64 - prefix)
ip.b = 0
} else {
ip.b = (ip.b >> (128 - prefix)) << (128 - prefix)
}
return ip
cidrSet *cidr.IpCidrSet
}
func (m *GeoIPMatcher) Init(cidrs []*CIDR) error {
ip4Count := 0
ip6Count := 0
for _, cidr := range cidrs {
ip := cidr.Ip
switch len(ip) {
case 4:
ip4Count++
case 16:
ip6Count++
default:
return fmt.Errorf("unexpect ip length: %d", len(ip))
}
}
cidrList := CIDRList(cidrs)
sort.Sort(&cidrList)
m.ip4 = make([]uint32, 0, ip4Count)
m.prefix4 = make([]uint8, 0, ip4Count)
m.ip6 = make([]ipv6, 0, ip6Count)
m.prefix6 = make([]uint8, 0, ip6Count)
for _, cidr := range cidrs {
ip := cidr.Ip
prefix := uint8(cidr.Prefix)
switch len(ip) {
case 4:
m.ip4 = append(m.ip4, normalize4(binary.BigEndian.Uint32(ip), prefix))
m.prefix4 = append(m.prefix4, prefix)
case 16:
ip6 := ipv6{
a: binary.BigEndian.Uint64(ip[0:8]),
b: binary.BigEndian.Uint64(ip[8:16]),
}
ip6 = normalize6(ip6, prefix)
m.ip6 = append(m.ip6, ip6)
m.prefix6 = append(m.prefix6, prefix)
addr, ok := netip.AddrFromSlice(cidr.Ip)
if !ok {
return fmt.Errorf("error when loading GeoIP: invalid IP: %s", cidr.Ip)
}
err := m.cidrSet.AddIpCidr(netip.PrefixFrom(addr, int(cidr.Prefix)))
if err != nil {
return fmt.Errorf("error when loading GeoIP: %w", err)
}
}
m.cidrSet.Merge()
return nil
}
@ -238,91 +146,13 @@ func (m *GeoIPMatcher) SetReverseMatch(isReverseMatch bool) {
m.reverseMatch = isReverseMatch
}
func (m *GeoIPMatcher) match4(ip uint32) bool {
if len(m.ip4) == 0 {
return false
}
if ip < m.ip4[0] {
return false
}
size := uint32(len(m.ip4))
l := uint32(0)
r := size
for l < r {
x := ((l + r) >> 1)
if ip < m.ip4[x] {
r = x
continue
}
nip := normalize4(ip, m.prefix4[x])
if nip == m.ip4[x] {
return true
}
l = x + 1
}
return l > 0 && normalize4(ip, m.prefix4[l-1]) == m.ip4[l-1]
}
func less6(a ipv6, b ipv6) bool {
return a.a < b.a || (a.a == b.a && a.b < b.b)
}
func (m *GeoIPMatcher) match6(ip ipv6) bool {
if len(m.ip6) == 0 {
return false
}
if less6(ip, m.ip6[0]) {
return false
}
size := uint32(len(m.ip6))
l := uint32(0)
r := size
for l < r {
x := (l + r) / 2
if less6(ip, m.ip6[x]) {
r = x
continue
}
if normalize6(ip, m.prefix6[x]) == m.ip6[x] {
return true
}
l = x + 1
}
return l > 0 && normalize6(ip, m.prefix6[l-1]) == m.ip6[l-1]
}
// Match returns true if the given ip is included by the GeoIP.
func (m *GeoIPMatcher) Match(ip net.IP) bool {
switch len(ip) {
case 4:
if m.reverseMatch {
return !m.match4(binary.BigEndian.Uint32(ip))
}
return m.match4(binary.BigEndian.Uint32(ip))
case 16:
if m.reverseMatch {
return !m.match6(ipv6{
a: binary.BigEndian.Uint64(ip[0:8]),
b: binary.BigEndian.Uint64(ip[8:16]),
})
}
return m.match6(ipv6{
a: binary.BigEndian.Uint64(ip[0:8]),
b: binary.BigEndian.Uint64(ip[8:16]),
})
default:
return false
func (m *GeoIPMatcher) Match(ip netip.Addr) bool {
match := m.cidrSet.IsContain(ip)
if m.reverseMatch {
return !match
}
return match
}
// GeoIPMatcherContainer is a container for GeoIPMatchers. It keeps unique copies of GeoIPMatcher by country code.
@ -344,6 +174,7 @@ func (c *GeoIPMatcherContainer) Add(geoip *GeoIP) (*GeoIPMatcher, error) {
m := &GeoIPMatcher{
countryCode: geoip.CountryCode,
reverseMatch: geoip.ReverseMatch,
cidrSet: cidr.NewIpCidrSet(),
}
if err := m.Init(geoip.Cidr); err != nil {
return nil, err
@ -369,8 +200,7 @@ func NewGeoIPMatcher(geoip *GeoIP) (*GeoIPMatcher, error) {
return matcher, nil
}
func (m *MultiGeoIPMatcher) ApplyIp(ip net.IP) bool {
func (m *MultiGeoIPMatcher) ApplyIp(ip netip.Addr) bool {
for _, matcher := range m.matchers {
if matcher.Match(ip) {
return true

View File

@ -41,7 +41,7 @@ func (gf *geoipFilter) Match(ip netip.Addr) bool {
return false
}
}
return !geoIPMatcher.Match(ip.AsSlice())
return !geoIPMatcher.Match(ip)
}
type ipnetFilter struct {

View File

@ -48,7 +48,7 @@ func (g *GEOIP) Match(metadata *C.Metadata) (bool, string) {
}
return false, g.adapter
}
return g.geoIPMatcher.Match(ip.AsSlice()), g.adapter
return g.geoIPMatcher.Match(ip), g.adapter
}
func (g *GEOIP) Adapter() string {