Skip to content

Commit 99b6ffa

Browse files
xishang0128xishang0128
committed
feat: add IP-ASN rule
1 parent 7ad37ca commit 99b6ffa

File tree

15 files changed

+248
-33
lines changed

15 files changed

+248
-33
lines changed

‎component/geodata/init.go

+30-3
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,11 @@ import (
1414
"github.com/metacubex/mihomo/log"
1515
)
1616

17-
var initGeoSite bool
18-
var initGeoIP int
17+
var (
18+
initGeoSite bool
19+
initGeoIP int
20+
initASN bool
21+
)
1922

2023
func InitGeoSite() error {
2124
if _, err := os.Stat(C.Path.GeoSite()); os.IsNotExist(err) {
@@ -113,7 +116,7 @@ func InitGeoIP() error {
113116
}
114117

115118
if initGeoIP != 2 {
116-
if !mmdb.Verify() {
119+
if !mmdb.Verify(C.Path.MMDB()) {
117120
log.Warnln("MMDB invalid, remove and download")
118121
if err := os.Remove(C.Path.MMDB()); err != nil {
119122
return fmt.Errorf("can't remove invalid MMDB: %s", err.Error())
@@ -126,3 +129,27 @@ func InitGeoIP() error {
126129
}
127130
return nil
128131
}
132+
133+
func InitASN() error {
134+
if _, err := os.Stat(C.Path.ASN()); os.IsNotExist(err) {
135+
log.Infoln("Can't find ASN.mmdb, start download")
136+
if err := mmdb.DownloadASN(C.Path.ASN()); err != nil {
137+
return fmt.Errorf("can't download ASN.mmdb: %s", err.Error())
138+
}
139+
log.Infoln("Download ASN.mmdb finish")
140+
initASN = false
141+
}
142+
if !initASN {
143+
if !mmdb.Verify(C.Path.ASN()) {
144+
log.Warnln("ASN invalid, remove and download")
145+
if err := os.Remove(C.Path.ASN()); err != nil {
146+
return fmt.Errorf("can't remove invalid ASN: %s", err.Error())
147+
}
148+
if err := mmdb.DownloadASN(C.Path.ASN()); err != nil {
149+
return fmt.Errorf("can't download ASN: %s", err.Error())
150+
}
151+
}
152+
initASN = true
153+
}
154+
return nil
155+
}

‎component/mmdb/mmdb.go

+57-18
Original file line numberDiff line numberDiff line change
@@ -25,56 +25,58 @@ const (
2525
)
2626

2727
var (
28-
reader Reader
29-
once sync.Once
28+
IPreader IPReader
29+
ASNreader ASNReader
30+
IPonce sync.Once
31+
ASNonce sync.Once
3032
)
3133

3234
func LoadFromBytes(buffer []byte) {
33-
once.Do(func() {
35+
IPonce.Do(func() {
3436
mmdb, err := maxminddb.FromBytes(buffer)
3537
if err != nil {
3638
log.Fatalln("Can't load mmdb: %s", err.Error())
3739
}
38-
reader = Reader{Reader: mmdb}
40+
IPreader = IPReader{Reader: mmdb}
3941
switch mmdb.Metadata.DatabaseType {
4042
case "sing-geoip":
41-
reader.databaseType = typeSing
43+
IPreader.databaseType = typeSing
4244
case "Meta-geoip0":
43-
reader.databaseType = typeMetaV0
45+
IPreader.databaseType = typeMetaV0
4446
default:
45-
reader.databaseType = typeMaxmind
47+
IPreader.databaseType = typeMaxmind
4648
}
4749
})
4850
}
4951

50-
func Verify() bool {
51-
instance, err := maxminddb.Open(C.Path.MMDB())
52+
func Verify(path string) bool {
53+
instance, err := maxminddb.Open(path)
5254
if err == nil {
5355
instance.Close()
5456
}
5557
return err == nil
5658
}
5759

58-
func Instance() Reader {
59-
once.Do(func() {
60+
func IPInstance() IPReader {
61+
IPonce.Do(func() {
6062
mmdbPath := C.Path.MMDB()
6163
log.Infoln("Load MMDB file: %s", mmdbPath)
6264
mmdb, err := maxminddb.Open(mmdbPath)
6365
if err != nil {
6466
log.Fatalln("Can't load MMDB: %s", err.Error())
6567
}
66-
reader = Reader{Reader: mmdb}
68+
IPreader = IPReader{Reader: mmdb}
6769
switch mmdb.Metadata.DatabaseType {
6870
case "sing-geoip":
69-
reader.databaseType = typeSing
71+
IPreader.databaseType = typeSing
7072
case "Meta-geoip0":
71-
reader.databaseType = typeMetaV0
73+
IPreader.databaseType = typeMetaV0
7274
default:
73-
reader.databaseType = typeMaxmind
75+
IPreader.databaseType = typeMaxmind
7476
}
7577
})
7678

77-
return reader
79+
return IPreader
7880
}
7981

8082
func DownloadMMDB(path string) (err error) {
@@ -96,6 +98,43 @@ func DownloadMMDB(path string) (err error) {
9698
return err
9799
}
98100

99-
func Reload() {
100-
mihomoOnce.Reset(&once)
101+
func ASNInstance() ASNReader {
102+
ASNonce.Do(func() {
103+
ASNPath := C.Path.ASN()
104+
log.Infoln("Load ASN file: %s", ASNPath)
105+
asn, err := maxminddb.Open(ASNPath)
106+
if err != nil {
107+
log.Fatalln("Can't load ASN: %s", err.Error())
108+
}
109+
ASNreader = ASNReader{Reader: asn}
110+
})
111+
112+
return ASNreader
113+
}
114+
115+
func DownloadASN(path string) (err error) {
116+
ctx, cancel := context.WithTimeout(context.Background(), time.Second*90)
117+
defer cancel()
118+
resp, err := mihomoHttp.HttpRequest(ctx, C.ASNUrl, http.MethodGet, http.Header{"User-Agent": {C.UA}}, nil)
119+
if err != nil {
120+
return
121+
}
122+
defer resp.Body.Close()
123+
124+
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644)
125+
if err != nil {
126+
return err
127+
}
128+
defer f.Close()
129+
_, err = io.Copy(f, resp.Body)
130+
131+
return err
132+
}
133+
134+
func ReloadIP() {
135+
mihomoOnce.Reset(&IPonce)
136+
}
137+
138+
func ReloadASN() {
139+
mihomoOnce.Reset(&ASNonce)
101140
}

‎component/mmdb/patch_android.go

+5-5
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,14 @@ package mmdb
55
import "github.com/oschwald/maxminddb-golang"
66

77
func InstallOverride(override *maxminddb.Reader) {
8-
newReader := Reader{Reader: override}
8+
newReader := IPReader{Reader: override}
99
switch override.Metadata.DatabaseType {
1010
case "sing-geoip":
11-
reader.databaseType = typeSing
11+
IPreader.databaseType = typeSing
1212
case "Meta-geoip0":
13-
reader.databaseType = typeMetaV0
13+
IPreader.databaseType = typeMetaV0
1414
default:
15-
reader.databaseType = typeMaxmind
15+
IPreader.databaseType = typeMaxmind
1616
}
17-
reader = newReader
17+
IPreader = newReader
1818
}

‎component/mmdb/reader.go

+17-2
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,21 @@ type geoip2Country struct {
1414
} `maxminddb:"country"`
1515
}
1616

17-
type Reader struct {
17+
type IPReader struct {
1818
*maxminddb.Reader
1919
databaseType
2020
}
2121

22-
func (r Reader) LookupCode(ipAddress net.IP) []string {
22+
type ASNReader struct {
23+
*maxminddb.Reader
24+
}
25+
26+
type ASNResult struct {
27+
AutonomousSystemNumber uint32 `maxminddb:"autonomous_system_number"`
28+
AutonomousSystemOrganization string `maxminddb:"autonomous_system_organization"`
29+
}
30+
31+
func (r IPReader) LookupCode(ipAddress net.IP) []string {
2332
switch r.databaseType {
2433
case typeMaxmind:
2534
var country geoip2Country
@@ -56,3 +65,9 @@ func (r Reader) LookupCode(ipAddress net.IP) []string {
5665
panic(fmt.Sprint("unknown geoip database type:", r.databaseType))
5766
}
5867
}
68+
69+
func (r ASNReader) LookupASN(ip net.IP) ASNResult {
70+
var result ASNResult
71+
r.Lookup(ip, &result)
72+
return result
73+
}

‎config/config.go

+3
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,7 @@ type RawConfig struct {
348348
type GeoXUrl struct {
349349
GeoIp string `yaml:"geoip" json:"geoip"`
350350
Mmdb string `yaml:"mmdb" json:"mmdb"`
351+
ASN string `yaml:"asn" json:"asn"`
351352
GeoSite string `yaml:"geosite" json:"geosite"`
352353
}
353354

@@ -495,6 +496,7 @@ func UnmarshalRawConfig(buf []byte) (*RawConfig, error) {
495496
},
496497
GeoXUrl: GeoXUrl{
497498
Mmdb: "https://github.com/MetaCubeX/meta-rules-dat/releases/download/latest/geoip.metadb",
499+
ASN: "https://github.com/P3TERX/GeoLite.mmdb/releases/download/2024.03.10/GeoLite2-ASN.mmdb",
498500
GeoIp: "https://github.com/MetaCubeX/meta-rules-dat/releases/download/latest/geoip.dat",
499501
GeoSite: "https://github.com/MetaCubeX/meta-rules-dat/releases/download/latest/geosite.dat",
500502
},
@@ -620,6 +622,7 @@ func parseGeneral(cfg *RawConfig) (*General, error) {
620622
C.GeoIpUrl = cfg.GeoXUrl.GeoIp
621623
C.GeoSiteUrl = cfg.GeoXUrl.GeoSite
622624
C.MmdbUrl = cfg.GeoXUrl.Mmdb
625+
C.ASNUrl = cfg.GeoXUrl.ASN
623626
C.GeodataMode = cfg.GeodataMode
624627
C.UA = cfg.GlobalUA
625628
if cfg.KeepAliveInterval != 0 {

‎config/update_geo.go

+21-2
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ func UpdateGeoDatabases() error {
3434
}
3535

3636
} else {
37-
defer mmdb.Reload()
37+
defer mmdb.ReloadIP()
3838
data, err := downloadForBytes(C.MmdbUrl)
3939
if err != nil {
4040
return fmt.Errorf("can't download MMDB database file: %w", err)
@@ -46,12 +46,31 @@ func UpdateGeoDatabases() error {
4646
}
4747
_ = instance.Close()
4848

49-
mmdb.Instance().Reader.Close() // mmdb is loaded with mmap, so it needs to be closed before overwriting the file
49+
mmdb.IPInstance().Reader.Close() // mmdb is loaded with mmap, so it needs to be closed before overwriting the file
5050
if err = saveFile(data, C.Path.MMDB()); err != nil {
5151
return fmt.Errorf("can't save MMDB database file: %w", err)
5252
}
5353
}
5454

55+
if C.ASNEnable {
56+
defer mmdb.ReloadASN()
57+
data, err := downloadForBytes(C.ASNUrl)
58+
if err != nil {
59+
return fmt.Errorf("can't download ASN database file: %w", err)
60+
}
61+
62+
instance, err := maxminddb.FromBytes(data)
63+
if err != nil {
64+
return fmt.Errorf("invalid ASN database file: %s", err)
65+
}
66+
_ = instance.Close()
67+
68+
mmdb.ASNInstance().Reader.Close()
69+
if err = saveFile(data, C.Path.ASN()); err != nil {
70+
return fmt.Errorf("can't save ASN database file: %w", err)
71+
}
72+
}
73+
5574
data, err := downloadForBytes(C.GeoSiteUrl)
5675
if err != nil {
5776
return fmt.Errorf("can't download GeoSite database file: %w", err)

‎constant/geodata.go

+2
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
package constant
22

33
var (
4+
ASNEnable bool
45
GeodataMode bool
56
GeoAutoUpdate bool
67
GeoUpdateInterval int
78
GeoIpUrl string
89
MmdbUrl string
910
GeoSiteUrl string
11+
ASNUrl string
1012
)

‎constant/metadata.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,8 @@ type Metadata struct {
133133
Type Type `json:"type"`
134134
SrcIP netip.Addr `json:"sourceIP"`
135135
DstIP netip.Addr `json:"destinationIP"`
136-
DstGeoIP []string `json:"destinationGeoIP"` // can be nil if never queried, empty slice if got no result
136+
DstGeoIP []string `json:"destinationGeoIP"` // can be nil if never queried, empty slice if got no result
137+
DstIPASN string `json:"destinationIPASN"`
137138
SrcPort uint16 `json:"sourcePort,string"` // `,string` is used to compatible with old version json output
138139
DstPort uint16 `json:"destinationPort,string"` // `,string` is used to compatible with old version json output
139140
InIP netip.Addr `json:"inboundIP"`

‎constant/path.go

+20
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ const Name = "mihomo"
1515
var (
1616
GeositeName = "GeoSite.dat"
1717
GeoipName = "GeoIP.dat"
18+
ASNName = "ASN.mmdb"
1819
)
1920

2021
// Path is used to get the configuration path
@@ -112,6 +113,25 @@ func (p *path) MMDB() string {
112113
return P.Join(p.homeDir, "geoip.metadb")
113114
}
114115

116+
func (p *path) ASN() string {
117+
files, err := os.ReadDir(p.homeDir)
118+
if err != nil {
119+
return ""
120+
}
121+
for _, fi := range files {
122+
if fi.IsDir() {
123+
// 目录则直接跳过
124+
continue
125+
} else {
126+
if strings.EqualFold(fi.Name(), "ASN.mmdb") {
127+
ASNName = fi.Name()
128+
return P.Join(p.homeDir, fi.Name())
129+
}
130+
}
131+
}
132+
return P.Join(p.homeDir, ASNName)
133+
}
134+
115135
func (p *path) OldCache() string {
116136
return P.Join(p.homeDir, ".cache")
117137
}

‎constant/rule.go

+3
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ const (
99
GEOSITE
1010
GEOIP
1111
IPCIDR
12+
IPASN
1213
SrcIPCIDR
1314
IPSuffix
1415
SrcIPSuffix
@@ -49,6 +50,8 @@ func (rt RuleType) String() string {
4950
return "GeoIP"
5051
case IPCIDR:
5152
return "IPCIDR"
53+
case IPASN:
54+
return "IPASN"
5255
case SrcIPCIDR:
5356
return "SrcIPCIDR"
5457
case IPSuffix:

‎dns/filters.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ var geoIPMatcher *router.GeoIPMatcher
2424

2525
func (gf *geoipFilter) Match(ip netip.Addr) bool {
2626
if !C.GeodataMode {
27-
codes := mmdb.Instance().LookupCode(ip.AsSlice())
27+
codes := mmdb.IPInstance().LookupCode(ip.AsSlice())
2828
for _, code := range codes {
2929
if !strings.EqualFold(code, gf.code) && !ip.IsPrivate() {
3030
return true

0 commit comments

Comments
 (0)