utils.go 4.3 KB

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