Skip to content

Commit

Permalink
chore: cleanup geo internal code
Browse files Browse the repository at this point in the history
  • Loading branch information
@wwqgtxx
wwqgtxx committed Sep 9, 2024
1 parent ef244b8 commit 7c8f451
Show file tree
Hide file tree
Showing 16 changed files with 218 additions and 213 deletions.
2 changes: 1 addition & 1 deletion adapter/provider/provider.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ func (pp *proxySetProvider) getSubscriptionInfo() {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
defer cancel()
resp, err := mihomoHttp.HttpRequestWithProxy(ctx, pp.Vehicle().(*resource.HTTPVehicle).Url(),
http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil, pp.Vehicle().Proxy())
http.MethodGet, nil, nil, pp.Vehicle().Proxy())
if err != nil {
return
}
Expand Down
141 changes: 95 additions & 46 deletions component/geodata/init.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,10 @@ import (
"io"
"net/http"
"os"
"sync"
"time"

"github.com/metacubex/mihomo/common/atomic"
mihomoHttp "github.com/metacubex/mihomo/component/http"
"github.com/metacubex/mihomo/component/mmdb"
C "github.com/metacubex/mihomo/constant"
Expand All @@ -18,55 +20,57 @@ var (
initGeoSite bool
initGeoIP int
initASN bool

initGeoSiteMutex sync.Mutex
initGeoIPMutex sync.Mutex
initASNMutex sync.Mutex

geoIpEnable atomic.Bool
geoSiteEnable atomic.Bool
asnEnable atomic.Bool

geoIpUrl string
mmdbUrl string
geoSiteUrl string
asnUrl string
)

func InitGeoSite() error {
if _, err := os.Stat(C.Path.GeoSite()); os.IsNotExist(err) {
log.Infoln("Can't find GeoSite.dat, start download")
if err := downloadGeoSite(C.Path.GeoSite()); err != nil {
return fmt.Errorf("can't download GeoSite.dat: %s", err.Error())
}
log.Infoln("Download GeoSite.dat finish")
initGeoSite = false
}
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 {
return fmt.Errorf("can't remove invalid GeoSite.dat: %s", err.Error())
}
if err := downloadGeoSite(C.Path.GeoSite()); err != nil {
return fmt.Errorf("can't download GeoSite.dat: %s", err.Error())
}
}
initGeoSite = true
}
return nil
func GeoIpUrl() string {
return geoIpUrl
}

func downloadGeoSite(path string) (err error) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
defer cancel()
resp, err := mihomoHttp.HttpRequest(ctx, C.GeoSiteUrl, http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil)
if err != nil {
return
}
defer resp.Body.Close()
func SetGeoIpUrl(url string) {
geoIpUrl = url
}

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)
func MmdbUrl() string {
return mmdbUrl
}

return err
func SetMmdbUrl(url string) {
mmdbUrl = url
}

func GeoSiteUrl() string {
return geoSiteUrl
}

func SetGeoSiteUrl(url string) {
geoSiteUrl = url
}

func downloadGeoIP(path string) (err error) {
func ASNUrl() string {
return asnUrl
}

func SetASNUrl(url string) {
asnUrl = url
}

func downloadToPath(url string, path string) (err error) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
defer cancel()
resp, err := mihomoHttp.HttpRequest(ctx, C.GeoIpUrl, http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil)
resp, err := mihomoHttp.HttpRequest(ctx, url, http.MethodGet, nil, nil)
if err != nil {
return
}
Expand All @@ -82,11 +86,41 @@ func downloadGeoIP(path string) (err error) {
return err
}

func InitGeoSite() error {
geoSiteEnable.Store(true)
initGeoSiteMutex.Lock()
defer initGeoSiteMutex.Unlock()
if _, err := os.Stat(C.Path.GeoSite()); os.IsNotExist(err) {
log.Infoln("Can't find GeoSite.dat, start download")
if err := downloadToPath(GeoSiteUrl(), C.Path.GeoSite()); err != nil {
return fmt.Errorf("can't download GeoSite.dat: %s", err.Error())
}
log.Infoln("Download GeoSite.dat finish")
initGeoSite = false
}
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 {
return fmt.Errorf("can't remove invalid GeoSite.dat: %s", err.Error())
}
if err := downloadToPath(GeoSiteUrl(), C.Path.GeoSite()); err != nil {
return fmt.Errorf("can't download GeoSite.dat: %s", err.Error())
}
}
initGeoSite = true
}
return nil
}

func InitGeoIP() error {
if C.GeodataMode {
geoIpEnable.Store(true)
initGeoIPMutex.Lock()
defer initGeoIPMutex.Unlock()
if 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 {
if err := downloadToPath(GeoIpUrl(), C.Path.GeoIP()); err != nil {
return fmt.Errorf("can't download GeoIP.dat: %s", err.Error())
}
log.Infoln("Download GeoIP.dat finish")
Expand All @@ -99,7 +133,7 @@ func InitGeoIP() error {
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 {
if err := downloadToPath(GeoIpUrl(), C.Path.GeoIP()); err != nil {
return fmt.Errorf("can't download GeoIP.dat: %s", err.Error())
}
}
Expand All @@ -110,7 +144,7 @@ func InitGeoIP() error {

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 {
if err := downloadToPath(MmdbUrl(), C.Path.MMDB()); err != nil {
return fmt.Errorf("can't download MMDB: %s", err.Error())
}
}
Expand All @@ -121,7 +155,7 @@ func InitGeoIP() error {
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 {
if err := downloadToPath(MmdbUrl(), C.Path.MMDB()); err != nil {
return fmt.Errorf("can't download MMDB: %s", err.Error())
}
}
Expand All @@ -131,9 +165,12 @@ func InitGeoIP() error {
}

func InitASN() error {
asnEnable.Store(true)
initASNMutex.Lock()
defer initASNMutex.Unlock()
if _, err := os.Stat(C.Path.ASN()); os.IsNotExist(err) {
log.Infoln("Can't find ASN.mmdb, start download")
if err := mmdb.DownloadASN(C.Path.ASN()); err != nil {
if err := downloadToPath(ASNUrl(), C.Path.ASN()); err != nil {
return fmt.Errorf("can't download ASN.mmdb: %s", err.Error())
}
log.Infoln("Download ASN.mmdb finish")
Expand All @@ -145,11 +182,23 @@ func InitASN() error {
if err := os.Remove(C.Path.ASN()); err != nil {
return fmt.Errorf("can't remove invalid ASN: %s", err.Error())
}
if err := mmdb.DownloadASN(C.Path.ASN()); err != nil {
if err := downloadToPath(ASNUrl(), C.Path.ASN()); err != nil {
return fmt.Errorf("can't download ASN: %s", err.Error())
}
}
initASN = true
}
return nil
}

func GeoIpEnable() bool {
return geoIpEnable.Load()
}

func GeoSiteEnable() bool {
return geoSiteEnable.Load()
}

func ASNEnable() bool {
return asnEnable.Load()
}
16 changes: 0 additions & 16 deletions component/geodata/utils.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,6 @@ import (

var (
geoMode bool
AutoUpdate bool
UpdateInterval int
geoLoaderName = "memconservative"
geoSiteMatcher = "succinct"
)
Expand All @@ -25,14 +23,6 @@ func GeodataMode() bool {
return geoMode
}

func GeoAutoUpdate() bool {
return AutoUpdate
}

func GeoUpdateInterval() int {
return UpdateInterval
}

func LoaderName() string {
return geoLoaderName
}
Expand All @@ -44,12 +34,6 @@ func SiteMatcherName() string {
func SetGeodataMode(newGeodataMode bool) {
geoMode = newGeodataMode
}
func SetGeoAutoUpdate(newAutoUpdate bool) {
AutoUpdate = newAutoUpdate
}
func SetGeoUpdateInterval(newGeoUpdateInterval int) {
UpdateInterval = newGeoUpdateInterval
}

func SetLoader(newLoader string) {
if newLoader == "memc" {
Expand Down
15 changes: 13 additions & 2 deletions component/http/http.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,21 @@ import (
"time"

"github.com/metacubex/mihomo/component/ca"
C "github.com/metacubex/mihomo/constant"
"github.com/metacubex/mihomo/listener/inner"
)

var (
ua string
)

func UA() string {
return ua
}

func SetUA(UA string) {
ua = UA
}

func HttpRequest(ctx context.Context, url, method string, header map[string][]string, body io.Reader) (*http.Response, error) {
return HttpRequestWithProxy(ctx, url, method, header, body, "")
}
Expand All @@ -35,7 +46,7 @@ func HttpRequestWithProxy(ctx context.Context, url, method string, header map[st
}

if _, ok := header["User-Agent"]; !ok {
req.Header.Set("User-Agent", C.UA)
req.Header.Set("User-Agent", UA())
}

if err != nil {
Expand Down
Loading

0 comments on commit 7c8f451

Please sign in to comment.