From e9a7e104c01614b50751257a77a62cee328c9968 Mon Sep 17 00:00:00 2001 From: metacubex Date: Mon, 9 Jan 2023 21:07:31 +0800 Subject: [PATCH] fix: geoip mmdb/geodata init --- component/geodata/init.go | 74 +++++++++++++++++++++++++++++++++-- component/mmdb/mmdb.go | 20 ++++++++++ config/initial.go | 82 +-------------------------------------- rules/common/base.go | 1 - rules/common/geoip.go | 5 +++ rules/common/geosite.go | 11 ++---- 6 files changed, 101 insertions(+), 92 deletions(-) diff --git a/component/geodata/init.go b/component/geodata/init.go index 2b36b626..f7dd7a9e 100644 --- a/component/geodata/init.go +++ b/component/geodata/init.go @@ -2,6 +2,7 @@ package geodata import ( "fmt" + "github.com/Dreamacro/clash/component/mmdb" C "github.com/Dreamacro/clash/constant" "github.com/Dreamacro/clash/log" "io" @@ -9,7 +10,8 @@ import ( "os" ) -var initFlag bool +var initGeoSite bool +var initGeoIP int func InitGeoSite() error { if _, err := os.Stat(C.Path.GeoSite()); os.IsNotExist(err) { @@ -18,8 +20,9 @@ func InitGeoSite() error { return fmt.Errorf("can't download GeoSite.dat: %s", err.Error()) } log.Infoln("Download GeoSite.dat finish") + initGeoSite = false } - if !initFlag { + if !initGeoSite { if err := Verify(C.GeositeName); err != nil { log.Warnln("GeoSite.dat invalid, remove and download: %s", err) if err := os.Remove(C.Path.GeoSite()); err != nil { @@ -29,7 +32,7 @@ func InitGeoSite() error { return fmt.Errorf("can't download GeoSite.dat: %s", err.Error()) } } - initFlag = true + initGeoSite = true } return nil } @@ -50,3 +53,68 @@ func downloadGeoSite(path string) (err error) { return err } + +func downloadGeoIP(path string) (err error) { + resp, err := http.Get(C.GeoIpUrl) + if err != nil { + return + } + defer resp.Body.Close() + + f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644) + if err != nil { + return err + } + defer f.Close() + _, err = io.Copy(f, resp.Body) + + return err +} + +func InitGeoIP() error { + if C.GeodataMode { + if _, err := os.Stat(C.Path.GeoIP()); os.IsNotExist(err) { + log.Infoln("Can't find GeoIP.dat, start download") + if err := downloadGeoIP(C.Path.GeoIP()); err != nil { + return fmt.Errorf("can't download GeoIP.dat: %s", err.Error()) + } + log.Infoln("Download GeoIP.dat finish") + initGeoIP = 0 + } + + if initGeoIP != 1 { + if err := Verify(C.GeoipName); err != nil { + log.Warnln("GeoIP.dat invalid, remove and download: %s", err) + if err := os.Remove(C.Path.GeoIP()); err != nil { + return fmt.Errorf("can't remove invalid GeoIP.dat: %s", err.Error()) + } + if err := downloadGeoIP(C.Path.GeoIP()); err != nil { + return fmt.Errorf("can't download GeoIP.dat: %s", err.Error()) + } + } + initGeoIP = 1 + } + return nil + } + + if _, err := os.Stat(C.Path.MMDB()); os.IsNotExist(err) { + log.Infoln("Can't find MMDB, start download") + if err := mmdb.DownloadMMDB(C.Path.MMDB()); err != nil { + return fmt.Errorf("can't download MMDB: %s", err.Error()) + } + } + + if initGeoIP != 2 { + if !mmdb.Verify() { + log.Warnln("MMDB invalid, remove and download") + if err := os.Remove(C.Path.MMDB()); err != nil { + return fmt.Errorf("can't remove invalid MMDB: %s", err.Error()) + } + if err := mmdb.DownloadMMDB(C.Path.MMDB()); err != nil { + return fmt.Errorf("can't download MMDB: %s", err.Error()) + } + } + initGeoIP = 2 + } + return nil +} diff --git a/component/mmdb/mmdb.go b/component/mmdb/mmdb.go index 58f3fe00..8f28d486 100644 --- a/component/mmdb/mmdb.go +++ b/component/mmdb/mmdb.go @@ -2,6 +2,9 @@ package mmdb import ( "github.com/oschwald/geoip2-golang" + "io" + "net/http" + "os" "sync" C "github.com/Dreamacro/clash/constant" @@ -42,3 +45,20 @@ func Instance() *geoip2.Reader { return mmdb } + +func DownloadMMDB(path string) (err error) { + resp, err := http.Get(C.MmdbUrl) + if err != nil { + return + } + defer resp.Body.Close() + + f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644) + if err != nil { + return err + } + defer f.Close() + _, err = io.Copy(f, resp.Body) + + return err +} diff --git a/config/initial.go b/config/initial.go index d819168f..0921040d 100644 --- a/config/initial.go +++ b/config/initial.go @@ -3,92 +3,12 @@ package config import ( "fmt" "github.com/Dreamacro/clash/component/geodata" - "github.com/Dreamacro/clash/component/mmdb" - "io" - "net/http" "os" C "github.com/Dreamacro/clash/constant" "github.com/Dreamacro/clash/log" ) -func downloadMMDB(path string) (err error) { - resp, err := http.Get(C.MmdbUrl) - if err != nil { - return - } - defer resp.Body.Close() - - f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644) - if err != nil { - return err - } - defer f.Close() - _, err = io.Copy(f, resp.Body) - - return err -} - -func downloadGeoIP(path string) (err error) { - resp, err := http.Get(C.GeoIpUrl) - if err != nil { - return - } - defer resp.Body.Close() - - f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644) - if err != nil { - return err - } - defer f.Close() - _, err = io.Copy(f, resp.Body) - - return err -} - -func initGeoIP() error { - if C.GeodataMode { - if _, err := os.Stat(C.Path.GeoIP()); os.IsNotExist(err) { - log.Infoln("Can't find GeoIP.dat, start download") - if err := downloadGeoIP(C.Path.GeoIP()); err != nil { - return fmt.Errorf("can't download GeoIP.dat: %s", err.Error()) - } - log.Infoln("Download GeoIP.dat finish") - } - - if err := geodata.Verify(C.GeoipName); err != nil { - log.Warnln("GeoIP.dat invalid, remove and download: %s", err) - if err := os.Remove(C.Path.GeoIP()); err != nil { - return fmt.Errorf("can't remove invalid GeoIP.dat: %s", err.Error()) - } - if err := downloadGeoIP(C.Path.GeoIP()); err != nil { - return fmt.Errorf("can't download GeoIP.dat: %s", err.Error()) - } - } - return nil - } - - if _, err := os.Stat(C.Path.MMDB()); os.IsNotExist(err) { - log.Infoln("Can't find MMDB, start download") - if err := downloadMMDB(C.Path.MMDB()); err != nil { - return fmt.Errorf("can't download MMDB: %s", err.Error()) - } - } - - if !mmdb.Verify() { - log.Warnln("MMDB invalid, remove and download") - if err := os.Remove(C.Path.MMDB()); err != nil { - return fmt.Errorf("can't remove invalid MMDB: %s", err.Error()) - } - - if err := downloadMMDB(C.Path.MMDB()); err != nil { - return fmt.Errorf("can't download MMDB: %s", err.Error()) - } - } - - return nil -} - // Init prepare necessary files func Init(dir string) error { // initial homedir @@ -122,7 +42,7 @@ func Init(dir string) error { C.GeoSiteUrl = rawCfg.GeoXUrl.GeoSite C.MmdbUrl = rawCfg.GeoXUrl.Mmdb // initial GeoIP - if err := initGeoIP(); err != nil { + if err := geodata.InitGeoIP(); err != nil { return fmt.Errorf("can't initial GeoIP: %w", err) } diff --git a/rules/common/base.go b/rules/common/base.go index abc2b434..d912107c 100644 --- a/rules/common/base.go +++ b/rules/common/base.go @@ -6,7 +6,6 @@ import ( var ( errPayload = errors.New("payloadRule error") - initFlag bool noResolve = "no-resolve" ) diff --git a/rules/common/geoip.go b/rules/common/geoip.go index 72e77045..0c134c63 100644 --- a/rules/common/geoip.go +++ b/rules/common/geoip.go @@ -71,6 +71,11 @@ func (g *GEOIP) GetRecodeSize() int { } func NewGEOIP(country string, adapter string, noResolveIP bool) (*GEOIP, error) { + if err := geodata.InitGeoIP(); err != nil { + log.Errorln("can't initial GeoIP: %s", err) + return nil, err + } + if !C.GeodataMode || strings.EqualFold(country, "LAN") { geoip := &GEOIP{ Base: &Base{}, diff --git a/rules/common/geosite.go b/rules/common/geosite.go index e5a0b9f7..865b0b1b 100644 --- a/rules/common/geosite.go +++ b/rules/common/geosite.go @@ -50,12 +50,9 @@ func (gs *GEOSITE) GetRecodeSize() int { } func NewGEOSITE(country string, adapter string) (*GEOSITE, error) { - if !initFlag { - if err := geodata.InitGeoSite(); err != nil { - log.Errorln("can't initial GeoSite: %s", err) - return nil, err - } - initFlag = true + if err := geodata.InitGeoSite(); err != nil { + log.Errorln("can't initial GeoSite: %s", err) + return nil, err } matcher, size, err := geodata.LoadGeoSiteMatcher(country) @@ -76,4 +73,4 @@ func NewGEOSITE(country string, adapter string) (*GEOSITE, error) { return geoSite, nil } -//var _ C.Rule = (*GEOSITE)(nil) +var _ C.Rule = (*GEOSITE)(nil)