mmdb.go 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  1. package mmdb
  2. import (
  3. "context"
  4. "io"
  5. "net/http"
  6. "os"
  7. "sync"
  8. "time"
  9. mihomoOnce "github.com/metacubex/mihomo/common/once"
  10. mihomoHttp "github.com/metacubex/mihomo/component/http"
  11. C "github.com/metacubex/mihomo/constant"
  12. "github.com/metacubex/mihomo/log"
  13. "github.com/oschwald/maxminddb-golang"
  14. )
  15. type databaseType = uint8
  16. const (
  17. typeMaxmind databaseType = iota
  18. typeSing
  19. typeMetaV0
  20. )
  21. var (
  22. IPreader IPReader
  23. ASNreader ASNReader
  24. IPonce sync.Once
  25. ASNonce sync.Once
  26. )
  27. func LoadFromBytes(buffer []byte) {
  28. IPonce.Do(func() {
  29. mmdb, err := maxminddb.FromBytes(buffer)
  30. if err != nil {
  31. log.Fatalln("Can't load mmdb: %s", err.Error())
  32. }
  33. IPreader = IPReader{Reader: mmdb}
  34. switch mmdb.Metadata.DatabaseType {
  35. case "sing-geoip":
  36. IPreader.databaseType = typeSing
  37. case "Meta-geoip0":
  38. IPreader.databaseType = typeMetaV0
  39. default:
  40. IPreader.databaseType = typeMaxmind
  41. }
  42. })
  43. }
  44. func Verify(path string) bool {
  45. instance, err := maxminddb.Open(path)
  46. if err == nil {
  47. instance.Close()
  48. }
  49. return err == nil
  50. }
  51. func IPInstance() IPReader {
  52. IPonce.Do(func() {
  53. mmdbPath := C.Path.MMDB()
  54. log.Infoln("Load MMDB file: %s", mmdbPath)
  55. mmdb, err := maxminddb.Open(mmdbPath)
  56. if err != nil {
  57. log.Fatalln("Can't load MMDB: %s", err.Error())
  58. }
  59. IPreader = IPReader{Reader: mmdb}
  60. switch mmdb.Metadata.DatabaseType {
  61. case "sing-geoip":
  62. IPreader.databaseType = typeSing
  63. case "Meta-geoip0":
  64. IPreader.databaseType = typeMetaV0
  65. default:
  66. IPreader.databaseType = typeMaxmind
  67. }
  68. })
  69. return IPreader
  70. }
  71. func DownloadMMDB(path string) (err error) {
  72. ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
  73. defer cancel()
  74. resp, err := mihomoHttp.HttpRequest(ctx, C.MmdbUrl, http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil)
  75. if err != nil {
  76. return
  77. }
  78. defer resp.Body.Close()
  79. f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644)
  80. if err != nil {
  81. return err
  82. }
  83. defer f.Close()
  84. _, err = io.Copy(f, resp.Body)
  85. return err
  86. }
  87. func ASNInstance() ASNReader {
  88. ASNonce.Do(func() {
  89. ASNPath := C.Path.ASN()
  90. log.Infoln("Load ASN file: %s", ASNPath)
  91. asn, err := maxminddb.Open(ASNPath)
  92. if err != nil {
  93. log.Fatalln("Can't load ASN: %s", err.Error())
  94. }
  95. ASNreader = ASNReader{Reader: asn}
  96. })
  97. return ASNreader
  98. }
  99. func DownloadASN(path string) (err error) {
  100. ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
  101. defer cancel()
  102. resp, err := mihomoHttp.HttpRequest(ctx, C.ASNUrl, http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil)
  103. if err != nil {
  104. return
  105. }
  106. defer resp.Body.Close()
  107. f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644)
  108. if err != nil {
  109. return err
  110. }
  111. defer f.Close()
  112. _, err = io.Copy(f, resp.Body)
  113. return err
  114. }
  115. func ReloadIP() {
  116. mihomoOnce.Reset(&IPonce)
  117. }
  118. func ReloadASN() {
  119. mihomoOnce.Reset(&ASNonce)
  120. }