chore: Add source matching for ip type rules

Author: xishang0128

constant/rule.go

@@ -8,8 +8,10 @@ const (
 	DomainRegex
 	GEOSITE
 	GEOIP
-	IPCIDR
+	SrcGEOIP
 	IPASN
+	SrcIPASN
+	IPCIDR
 	SrcIPCIDR
 	IPSuffix
 	SrcIPSuffix
@@ -48,10 +50,14 @@ func (rt RuleType) String() string {
 		return "GeoSite"
 	case GEOIP:
 		return "GeoIP"
-	case IPCIDR:
-		return "IPCIDR"
+	case SrcGEOIP:
+		return "SrcGeoIP"
 	case IPASN:
 		return "IPASN"
+	case SrcIPASN:
+		return "SrcIPASN"
+	case IPCIDR:
+		return "IPCIDR"
 	case SrcIPCIDR:
 		return "SrcIPCIDR"
 	case IPSuffix:

rules/common/geoip.go

@@ -17,18 +17,25 @@ type GEOIP struct {
 	country      string
 	adapter      string
 	noResolveIP  bool
+	isSourceIP   bool
 	geoIPMatcher *router.GeoIPMatcher
 	recodeSize   int
 }
 
 var _ C.Rule = (*GEOIP)(nil)
 
 func (g *GEOIP) RuleType() C.RuleType {
+	if g.isSourceIP {
+		return C.SrcGEOIP
+	}
 	return C.GEOIP
 }
 
 func (g *GEOIP) Match(metadata *C.Metadata) (bool, string) {
 	ip := metadata.DstIP
+	if g.isSourceIP {
+		ip = metadata.SrcIP
+	}
 	if !ip.IsValid() {
 		return false, ""
 	}
@@ -49,6 +56,16 @@ func (g *GEOIP) Match(metadata *C.Metadata) (bool, string) {
 	}
 
 	if !C.GeodataMode {
+		if g.isSourceIP {
+			codes := mmdb.IPInstance().LookupCode(ip.AsSlice())
+			for _, code := range codes {
+				if g.country == code {
+					return true, g.adapter
+				}
+			}
+			return false, g.adapter
+		}
+
 		if metadata.DstGeoIP != nil {
 			return false, g.adapter
 		}
@@ -62,7 +79,7 @@ func (g *GEOIP) Match(metadata *C.Metadata) (bool, string) {
 	}
 
 	match := g.geoIPMatcher.Match(ip)
-	if match {
+	if match && !g.isSourceIP {
 		metadata.DstGeoIP = append(metadata.DstGeoIP, g.country)
 	}
 	return match, g.adapter
@@ -92,7 +109,7 @@ func (g *GEOIP) GetRecodeSize() int {
 	return g.recodeSize
 }
 
-func NewGEOIP(country string, adapter string, noResolveIP bool) (*GEOIP, error) {
+func NewGEOIP(country string, adapter string, isSrc, noResolveIP bool) (*GEOIP, error) {
 	if err := geodata.InitGeoIP(); err != nil {
 		log.Errorln("can't initial GeoIP: %s", err)
 		return nil, err
@@ -105,6 +122,7 @@ func NewGEOIP(country string, adapter string, noResolveIP bool) (*GEOIP, error)
 			country:     country,
 			adapter:     adapter,
 			noResolveIP: noResolveIP,
+			isSourceIP:  isSrc,
 		}
 		return geoip, nil
 	}
@@ -120,6 +138,7 @@ func NewGEOIP(country string, adapter string, noResolveIP bool) (*GEOIP, error)
 		country:      country,
 		adapter:      adapter,
 		noResolveIP:  noResolveIP,
+		isSourceIP:   isSrc,
 		geoIPMatcher: geoIPMatcher,
 		recodeSize:   size,
 	}

rules/common/ipasn.go

@@ -14,24 +14,32 @@ type ASN struct {
 	asn         string
 	adapter     string
 	noResolveIP bool
+	isSourceIP  bool
 }
 
 func (a *ASN) Match(metadata *C.Metadata) (bool, string) {
 	ip := metadata.DstIP
+	if a.isSourceIP {
+		ip = metadata.SrcIP
+	}
 	if !ip.IsValid() {
 		return false, ""
 	}
 
 	result := mmdb.ASNInstance().LookupASN(ip.AsSlice())
-
 	asnNumber := strconv.FormatUint(uint64(result.AutonomousSystemNumber), 10)
-	metadata.DstIPASN = asnNumber + " " + result.AutonomousSystemOrganization
+	if !a.isSourceIP {
+		metadata.DstIPASN = asnNumber + " " + result.AutonomousSystemOrganization
+	}
 
 	match := a.asn == asnNumber
 	return match, a.adapter
 }
 
 func (a *ASN) RuleType() C.RuleType {
+	if a.isSourceIP {
+		return C.SrcIPASN
+	}
 	return C.IPASN
 }
 
@@ -51,7 +59,7 @@ func (a *ASN) GetASN() string {
 	return a.asn
 }
 
-func NewIPASN(asn string, adapter string, noResolveIP bool) (*ASN, error) {
+func NewIPASN(asn string, adapter string, isSrc, noResolveIP bool) (*ASN, error) {
 	C.ASNEnable = true
 	if err := geodata.InitASN(); err != nil {
 		log.Errorln("can't initial ASN: %s", err)
@@ -63,5 +71,6 @@ func NewIPASN(asn string, adapter string, noResolveIP bool) (*ASN, error) {
 		asn:         asn,
 		adapter:     adapter,
 		noResolveIP: noResolveIP,
+		isSourceIP:  isSrc,
 	}, nil
 }

rules/parser.go

@@ -23,13 +23,17 @@ func ParseRule(tp, payload, target string, params []string, subRules map[string]
 		parsed, parseErr = RC.NewGEOSITE(payload, target)
 	case "GEOIP":
 		noResolve := RC.HasNoResolve(params)
-		parsed, parseErr = RC.NewGEOIP(payload, target, noResolve)
+		parsed, parseErr = RC.NewGEOIP(payload, target, false, noResolve)
+	case "SRC-GEOIP":
+		parsed, parseErr = RC.NewGEOIP(payload, target, true, true)
+	case "IP-ASN":
+		noResolve := RC.HasNoResolve(params)
+		parsed, parseErr = RC.NewIPASN(payload, target, false, noResolve)
+	case "SRC-IP-ASN":
+		parsed, parseErr = RC.NewIPASN(payload, target, true, true)
 	case "IP-CIDR", "IP-CIDR6":
 		noResolve := RC.HasNoResolve(params)
 		parsed, parseErr = RC.NewIPCIDR(payload, target, RC.WithIPCIDRNoResolve(noResolve))
-	case "IP-ASN":
-		noResolve := RC.HasNoResolve(params)
-		parsed, parseErr = RC.NewIPASN(payload, target, noResolve)
 	case "SRC-IP-CIDR":
 		parsed, parseErr = RC.NewIPCIDR(payload, target, RC.WithIPCIDRSourceIP(true), RC.WithIPCIDRNoResolve(true))
 	case "IP-SUFFIX":