cache.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. package memconservative
  2. import (
  3. "fmt"
  4. "os"
  5. "strings"
  6. "github.com/metacubex/mihomo/component/geodata/router"
  7. C "github.com/metacubex/mihomo/constant"
  8. "github.com/metacubex/mihomo/log"
  9. "google.golang.org/protobuf/proto"
  10. )
  11. type GeoIPCache map[string]*router.GeoIP
  12. func (g GeoIPCache) Has(key string) bool {
  13. return !(g.Get(key) == nil)
  14. }
  15. func (g GeoIPCache) Get(key string) *router.GeoIP {
  16. if g == nil {
  17. return nil
  18. }
  19. return g[key]
  20. }
  21. func (g GeoIPCache) Set(key string, value *router.GeoIP) {
  22. if g == nil {
  23. g = make(map[string]*router.GeoIP)
  24. }
  25. g[key] = value
  26. }
  27. func (g GeoIPCache) Unmarshal(filename, code string) (*router.GeoIP, error) {
  28. asset := C.Path.GetAssetLocation(filename)
  29. idx := strings.ToLower(asset + ":" + code)
  30. if g.Has(idx) {
  31. return g.Get(idx), nil
  32. }
  33. geoipBytes, err := Decode(asset, code)
  34. switch err {
  35. case nil:
  36. var geoip router.GeoIP
  37. if err := proto.Unmarshal(geoipBytes, &geoip); err != nil {
  38. return nil, err
  39. }
  40. g.Set(idx, &geoip)
  41. return &geoip, nil
  42. case errCodeNotFound:
  43. return nil, fmt.Errorf("country code %s%s%s", code, " not found in ", filename)
  44. case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
  45. errInvalidGeodataFile, errInvalidGeodataVarintLength:
  46. log.Warnln("failed to decode geoip file: %s%s", filename, ", fallback to the original ReadFile method")
  47. geoipBytes, err = os.ReadFile(asset)
  48. if err != nil {
  49. return nil, err
  50. }
  51. var geoipList router.GeoIPList
  52. if err := proto.Unmarshal(geoipBytes, &geoipList); err != nil {
  53. return nil, err
  54. }
  55. for _, geoip := range geoipList.GetEntry() {
  56. if strings.EqualFold(code, geoip.GetCountryCode()) {
  57. g.Set(idx, geoip)
  58. return geoip, nil
  59. }
  60. }
  61. default:
  62. return nil, err
  63. }
  64. return nil, fmt.Errorf("country code %s%s%s", code, " not found in ", filename)
  65. }
  66. type GeoSiteCache map[string]*router.GeoSite
  67. func (g GeoSiteCache) Has(key string) bool {
  68. return !(g.Get(key) == nil)
  69. }
  70. func (g GeoSiteCache) Get(key string) *router.GeoSite {
  71. if g == nil {
  72. return nil
  73. }
  74. return g[key]
  75. }
  76. func (g GeoSiteCache) Set(key string, value *router.GeoSite) {
  77. if g == nil {
  78. g = make(map[string]*router.GeoSite)
  79. }
  80. g[key] = value
  81. }
  82. func (g GeoSiteCache) Unmarshal(filename, code string) (*router.GeoSite, error) {
  83. asset := C.Path.GetAssetLocation(filename)
  84. idx := strings.ToLower(asset + ":" + code)
  85. if g.Has(idx) {
  86. return g.Get(idx), nil
  87. }
  88. geositeBytes, err := Decode(asset, code)
  89. switch err {
  90. case nil:
  91. var geosite router.GeoSite
  92. if err := proto.Unmarshal(geositeBytes, &geosite); err != nil {
  93. return nil, err
  94. }
  95. g.Set(idx, &geosite)
  96. return &geosite, nil
  97. case errCodeNotFound:
  98. return nil, fmt.Errorf("list %s%s%s", code, " not found in ", filename)
  99. case errFailedToReadBytes, errFailedToReadExpectedLenBytes,
  100. errInvalidGeodataFile, errInvalidGeodataVarintLength:
  101. log.Warnln("failed to decode geosite file: %s%s", filename, ", fallback to the original ReadFile method")
  102. geositeBytes, err = os.ReadFile(asset)
  103. if err != nil {
  104. return nil, err
  105. }
  106. var geositeList router.GeoSiteList
  107. if err := proto.Unmarshal(geositeBytes, &geositeList); err != nil {
  108. return nil, err
  109. }
  110. for _, geosite := range geositeList.GetEntry() {
  111. if strings.EqualFold(code, geosite.GetCountryCode()) {
  112. g.Set(idx, geosite)
  113. return geosite, nil
  114. }
  115. }
  116. default:
  117. return nil, err
  118. }
  119. return nil, fmt.Errorf("list %s%s%s", code, " not found in ", filename)
  120. }