chore: Add GeoIP result to metadata

MetaCubeX · Mar 2, 2024 · d273408 · d273408
1 parent 7eb16a0
commit d273408
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 13 deletions.
diff --git a/component/mmdb/reader.go b/component/mmdb/reader.go
@@ -5,7 +5,6 @@ import (
 	"net"
 
 	"github.com/oschwald/maxminddb-golang"
-	"github.com/sagernet/sing/common"
 )
 
 type geoip2Country struct {
@@ -44,9 +43,11 @@ func (r Reader) LookupCode(ipAddress net.IP) []string {
 		case string:
 			return []string{record}
 		case []any: // lookup returned type of slice is []any
-			return common.Map(record, func(it any) string {
-				return it.(string)
-			})
+			result := make([]string, 0, len(record))
+			for _, item := range record {
+				result = append(result, item.(string))
+			}
+			return result
 		}
 		return []string{}
 

diff --git a/constant/metadata.go b/constant/metadata.go
@@ -133,6 +133,7 @@ type Metadata struct {
 	Type         Type       `json:"type"`
 	SrcIP        netip.Addr `json:"sourceIP"`
 	DstIP        netip.Addr `json:"destinationIP"`
+	DstGeoIP     []string   `json:"destinationGeoIP"`       // can be nil if never queried, empty slice if got no result
 	SrcPort      uint16     `json:"sourcePort,string"`      // `,string` is used to compatible with old version json output
 	DstPort      uint16     `json:"destinationPort,string"` // `,string` is used to compatible with old version json output
 	InIP         netip.Addr `json:"inboundIP"`

diff --git a/rules/common/geoip.go b/rules/common/geoip.go
@@ -21,6 +21,8 @@ type GEOIP struct {
 	recodeSize   int
 }
 
+var _ C.Rule = (*GEOIP)(nil)
+
 func (g *GEOIP) RuleType() C.RuleType {
 	return C.GEOIP
 }
@@ -31,24 +33,39 @@ func (g *GEOIP) Match(metadata *C.Metadata) (bool, string) {
 		return false, ""
 	}
 
-	if strings.EqualFold(g.country, "LAN") {
+	if g.country == "lan" {
 		return ip.IsPrivate() ||
 			ip.IsUnspecified() ||
 			ip.IsLoopback() ||
 			ip.IsMulticast() ||
 			ip.IsLinkLocalUnicast() ||
 			resolver.IsFakeBroadcastIP(ip), g.adapter
 	}
+
+	for _, code := range metadata.DstGeoIP {
+		if g.country == code {
+			return true, g.adapter
+		}
+	}
+
 	if !C.GeodataMode {
-		codes := mmdb.Instance().LookupCode(ip.AsSlice())
-		for _, code := range codes {
-			if strings.EqualFold(code, g.country) {
+		if metadata.DstGeoIP != nil {
+			return false, g.adapter
+		}
+		metadata.DstGeoIP = mmdb.Instance().LookupCode(ip.AsSlice())
+		for _, code := range metadata.DstGeoIP {
+			if g.country == code {
 				return true, g.adapter
 			}
 		}
 		return false, g.adapter
 	}
-	return g.geoIPMatcher.Match(ip), g.adapter
+
+	match := g.geoIPMatcher.Match(ip)
+	if match {
+		metadata.DstGeoIP = append(metadata.DstGeoIP, g.country)
+	}
+	return match, g.adapter
 }
 
 func (g *GEOIP) Adapter() string {
@@ -80,8 +97,9 @@ func NewGEOIP(country string, adapter string, noResolveIP bool) (*GEOIP, error)
 		log.Errorln("can't initial GeoIP: %s", err)
 		return nil, err
 	}
+	country = strings.ToLower(country)
 
-	if !C.GeodataMode || strings.EqualFold(country, "LAN") {
+	if !C.GeodataMode || country == "lan" {
 		geoip := &GEOIP{
 			Base:        &Base{},
 			country:     country,
@@ -93,7 +111,7 @@ func NewGEOIP(country string, adapter string, noResolveIP bool) (*GEOIP, error)
 
 	geoIPMatcher, size, err := geodata.LoadGeoIPMatcher(country)
 	if err != nil {
-		return nil, fmt.Errorf("[GeoIP] %s", err.Error())
+		return nil, fmt.Errorf("[GeoIP] %w", err)
 	}
 
 	log.Infoln("Start initial GeoIP rule %s => %s, records: %d", country, adapter, size)
@@ -107,5 +125,3 @@ func NewGEOIP(country string, adapter string, noResolveIP bool) (*GEOIP, error)
 	}
 	return geoip, nil
 }
-
-//var _ C.Rule = (*GEOIP)(nil)