chore: Increase support for other format of ASN

MetaCubeX · Nov 4, 2024 · a86c562 · a86c562
1 parent 3e966e8
commit a86c562
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 14 deletions.
diff --git a/component/mmdb/reader.go b/component/mmdb/reader.go
@@ -5,6 +5,7 @@ import (
 	"net"
 	"strings"
 
+	"github.com/metacubex/mihomo/log"
 	"github.com/oschwald/maxminddb-golang"
 )
 
@@ -23,11 +24,16 @@ type ASNReader struct {
 	*maxminddb.Reader
 }
 
-type ASNResult struct {
+type GeoLite2 struct {
 	AutonomousSystemNumber       uint32 `maxminddb:"autonomous_system_number"`
 	AutonomousSystemOrganization string `maxminddb:"autonomous_system_organization"`
 }
 
+type IPInfo struct {
+	ASN  string `maxminddb:"asn"`
+	Name string `maxminddb:"name"`
+}
+
 func (r IPReader) LookupCode(ipAddress net.IP) []string {
 	switch r.databaseType {
 	case typeMaxmind:
@@ -66,8 +72,18 @@ func (r IPReader) LookupCode(ipAddress net.IP) []string {
 	}
 }
 
-func (r ASNReader) LookupASN(ip net.IP) ASNResult {
-	var result ASNResult
-	r.Lookup(ip, &result)
-	return result
+func (r ASNReader) LookupASN(ip net.IP) (string, string) {
+	switch r.Metadata.DatabaseType {
+	case "GeoLite2-ASN", "DBIP-ASN-Lite (compat=GeoLite2-ASN)":
+		var result GeoLite2
+		_ = r.Lookup(ip, &result)
+		return fmt.Sprint(result.AutonomousSystemNumber), result.AutonomousSystemOrganization
+	case "ipinfo generic_asn_free.mmdb":
+		var result IPInfo
+		_ = r.Lookup(ip, &result)
+		return result.ASN[2:], result.Name
+	default:
+		log.Warnln("Unsupported ASN type: %s", r.Metadata.DatabaseType)
+	}
+	return "", ""
 }
diff --git a/rules/common/ipasn.go b/rules/common/ipasn.go
@@ -1,8 +1,6 @@
 package common
 
 import (
-	"strconv"
-
 	"github.com/metacubex/mihomo/component/geodata"
 	"github.com/metacubex/mihomo/component/mmdb"
 	C "github.com/metacubex/mihomo/constant"
@@ -26,17 +24,14 @@ func (a *ASN) Match(metadata *C.Metadata) (bool, string) {
 		return false, ""
 	}
 
-	result := mmdb.ASNInstance().LookupASN(ip.AsSlice())
-	asnNumber := strconv.FormatUint(uint64(result.AutonomousSystemNumber), 10)
-	ipASN := asnNumber + " " + result.AutonomousSystemOrganization
+	asn, aso := mmdb.ASNInstance().LookupASN(ip.AsSlice())
 	if a.isSourceIP {
-		metadata.SrcIPASN = ipASN
+		metadata.SrcIPASN = asn + " " + aso
 	} else {
-		metadata.DstIPASN = ipASN
+		metadata.DstIPASN = asn + " " + aso
 	}
 
-	match := a.asn == asnNumber
-	return match, a.adapter
+	return a.asn == asn, a.adapter
 }
 
 func (a *ASN) RuleType() C.RuleType {