github.com/metacubex/mihomo@v1.18.5/component/geodata/utils.go (about)

     1  package geodata
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"strings"
     7  
     8  	"golang.org/x/sync/singleflight"
     9  
    10  	"github.com/metacubex/mihomo/component/geodata/router"
    11  	C "github.com/metacubex/mihomo/constant"
    12  	"github.com/metacubex/mihomo/log"
    13  )
    14  
    15  var (
    16  	geoMode        bool
    17  	AutoUpdate     bool
    18  	UpdateInterval int
    19  	geoLoaderName  = "memconservative"
    20  	geoSiteMatcher = "succinct"
    21  )
    22  
    23  //  geoLoaderName = "standard"
    24  
    25  func GeodataMode() bool {
    26  	return geoMode
    27  }
    28  
    29  func GeoAutoUpdate() bool {
    30  	return AutoUpdate
    31  }
    32  
    33  func GeoUpdateInterval() int {
    34  	return UpdateInterval
    35  }
    36  
    37  func LoaderName() string {
    38  	return geoLoaderName
    39  }
    40  
    41  func SiteMatcherName() string {
    42  	return geoSiteMatcher
    43  }
    44  
    45  func SetGeodataMode(newGeodataMode bool) {
    46  	geoMode = newGeodataMode
    47  }
    48  func SetGeoAutoUpdate(newAutoUpdate bool) {
    49  	AutoUpdate = newAutoUpdate
    50  }
    51  func SetGeoUpdateInterval(newGeoUpdateInterval int) {
    52  	UpdateInterval = newGeoUpdateInterval
    53  }
    54  
    55  func SetLoader(newLoader string) {
    56  	if newLoader == "memc" {
    57  		newLoader = "memconservative"
    58  	}
    59  	geoLoaderName = newLoader
    60  }
    61  
    62  func SetSiteMatcher(newMatcher string) {
    63  	switch newMatcher {
    64  	case "mph", "hybrid":
    65  		geoSiteMatcher = "mph"
    66  	default:
    67  		geoSiteMatcher = "succinct"
    68  	}
    69  }
    70  
    71  func Verify(name string) error {
    72  	switch name {
    73  	case C.GeositeName:
    74  		_, _, err := LoadGeoSiteMatcher("CN")
    75  		return err
    76  	case C.GeoipName:
    77  		_, _, err := LoadGeoIPMatcher("CN")
    78  		return err
    79  	default:
    80  		return fmt.Errorf("not support name")
    81  	}
    82  }
    83  
    84  var loadGeoSiteMatcherSF = singleflight.Group{}
    85  
    86  func LoadGeoSiteMatcher(countryCode string) (router.DomainMatcher, int, error) {
    87  	if countryCode == "" {
    88  		return nil, 0, fmt.Errorf("country code could not be empty")
    89  	}
    90  
    91  	not := false
    92  	if countryCode[0] == '!' {
    93  		not = true
    94  		countryCode = countryCode[1:]
    95  	}
    96  	countryCode = strings.ToLower(countryCode)
    97  
    98  	parts := strings.Split(countryCode, "@")
    99  	if len(parts) == 0 {
   100  		return nil, 0, errors.New("empty rule")
   101  	}
   102  	listName := strings.TrimSpace(parts[0])
   103  	attrVal := parts[1:]
   104  
   105  	if listName == "" {
   106  		return nil, 0, fmt.Errorf("empty listname in rule: %s", countryCode)
   107  	}
   108  
   109  	v, err, shared := loadGeoSiteMatcherSF.Do(listName, func() (interface{}, error) {
   110  		geoLoader, err := GetGeoDataLoader(geoLoaderName)
   111  		if err != nil {
   112  			return nil, err
   113  		}
   114  		return geoLoader.LoadGeoSite(listName)
   115  	})
   116  	if err != nil {
   117  		if !shared {
   118  			loadGeoSiteMatcherSF.Forget(listName) // don't store the error result
   119  		}
   120  		return nil, 0, err
   121  	}
   122  	domains := v.([]*router.Domain)
   123  
   124  	attrs := parseAttrs(attrVal)
   125  	if attrs.IsEmpty() {
   126  		if strings.Contains(countryCode, "@") {
   127  			log.Warnln("empty attribute list: %s", countryCode)
   128  		}
   129  	} else {
   130  		filteredDomains := make([]*router.Domain, 0, len(domains))
   131  		hasAttrMatched := false
   132  		for _, domain := range domains {
   133  			if attrs.Match(domain) {
   134  				hasAttrMatched = true
   135  				filteredDomains = append(filteredDomains, domain)
   136  			}
   137  		}
   138  		if !hasAttrMatched {
   139  			log.Warnln("attribute match no rule: geosite: %s", countryCode)
   140  		}
   141  		domains = filteredDomains
   142  	}
   143  
   144  	/**
   145  	linear: linear algorithm
   146  	matcher, err := router.NewDomainMatcher(domains)
   147  	mph:minimal perfect hash algorithm
   148  	*/
   149  	var matcher router.DomainMatcher
   150  	if geoSiteMatcher == "mph" {
   151  		matcher, err = router.NewMphMatcherGroup(domains, not)
   152  	} else {
   153  		matcher, err = router.NewSuccinctMatcherGroup(domains, not)
   154  	}
   155  	if err != nil {
   156  		return nil, 0, err
   157  	}
   158  
   159  	return matcher, len(domains), nil
   160  }
   161  
   162  var loadGeoIPMatcherSF = singleflight.Group{}
   163  
   164  func LoadGeoIPMatcher(country string) (*router.GeoIPMatcher, int, error) {
   165  	if len(country) == 0 {
   166  		return nil, 0, fmt.Errorf("country code could not be empty")
   167  	}
   168  
   169  	not := false
   170  	if country[0] == '!' {
   171  		not = true
   172  		country = country[1:]
   173  	}
   174  	country = strings.ToLower(country)
   175  
   176  	v, err, shared := loadGeoIPMatcherSF.Do(country, func() (interface{}, error) {
   177  		geoLoader, err := GetGeoDataLoader(geoLoaderName)
   178  		if err != nil {
   179  			return nil, err
   180  		}
   181  		return geoLoader.LoadGeoIP(country)
   182  	})
   183  	if err != nil {
   184  		if !shared {
   185  			loadGeoIPMatcherSF.Forget(country) // don't store the error result
   186  		}
   187  		return nil, 0, err
   188  	}
   189  	records := v.([]*router.CIDR)
   190  
   191  	geoIP := &router.GeoIP{
   192  		CountryCode:  country,
   193  		Cidr:         records,
   194  		ReverseMatch: not,
   195  	}
   196  
   197  	matcher, err := router.NewGeoIPMatcher(geoIP)
   198  	if err != nil {
   199  		return nil, 0, err
   200  	}
   201  	return matcher, len(records), nil
   202  }
   203  
   204  func ClearCache() {
   205  	loadGeoSiteMatcherSF = singleflight.Group{}
   206  	loadGeoIPMatcherSF = singleflight.Group{}
   207  }