From af42ed00a28f9d8b851377ec98476303efd04907 Mon Sep 17 00:00:00 2001 From: xishang0128 Date: Mon, 1 Apr 2024 17:59:18 +0800 Subject: [PATCH] chore: Add source matching for ip type rules --- constant/rule.go | 12 +++++++++--- rules/common/geoip.go | 39 +++++++++++++++++++++++++-------------- rules/common/ipasn.go | 15 ++++++++++++--- rules/parser.go | 12 ++++++++---- 4 files changed, 54 insertions(+), 24 deletions(-) diff --git a/constant/rule.go b/constant/rule.go index fcefaba6ca..8fdd75a6c0 100644 --- a/constant/rule.go +++ b/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: diff --git a/rules/common/geoip.go b/rules/common/geoip.go index 223a79042d..2d130a9953 100644 --- a/rules/common/geoip.go +++ b/rules/common/geoip.go @@ -17,6 +17,7 @@ type GEOIP struct { country string adapter string noResolveIP bool + isSourceIP bool geoIPMatcher *router.GeoIPMatcher recodeSize int } @@ -24,11 +25,17 @@ type GEOIP struct { 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, "" } @@ -42,27 +49,29 @@ func (g *GEOIP) Match(metadata *C.Metadata) (bool, string) { resolver.IsFakeBroadcastIP(ip), g.adapter } - for _, code := range metadata.DstGeoIP { - if g.country == code { - return true, g.adapter - } - } - - if !C.GeodataMode { - if metadata.DstGeoIP != nil { - return false, g.adapter - } - metadata.DstGeoIP = mmdb.IPInstance().LookupCode(ip.AsSlice()) + if !g.isSourceIP { for _, code := range metadata.DstGeoIP { if g.country == code { return true, g.adapter } } - return false, g.adapter + + if !C.GeodataMode { + if metadata.DstGeoIP != nil { + return false, g.adapter + } + metadata.DstGeoIP = mmdb.IPInstance().LookupCode(ip.AsSlice()) + for _, code := range metadata.DstGeoIP { + if g.country == code { + return true, g.adapter + } + } + return false, g.adapter + } } match := g.geoIPMatcher.Match(ip) - if match { + if match && !g.isSourceIP { metadata.DstGeoIP = append(metadata.DstGeoIP, g.country) } return match, g.adapter @@ -92,7 +101,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 +114,7 @@ func NewGEOIP(country string, adapter string, noResolveIP bool) (*GEOIP, error) country: country, adapter: adapter, noResolveIP: noResolveIP, + isSourceIP: isSrc, } return geoip, nil } @@ -120,6 +130,7 @@ func NewGEOIP(country string, adapter string, noResolveIP bool) (*GEOIP, error) country: country, adapter: adapter, noResolveIP: noResolveIP, + isSourceIP: isSrc, geoIPMatcher: geoIPMatcher, recodeSize: size, } diff --git a/rules/common/ipasn.go b/rules/common/ipasn.go index 1fce8af482..df4b6531c6 100644 --- a/rules/common/ipasn.go +++ b/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 } diff --git a/rules/parser.go b/rules/parser.go index b69cc4fd9c..032a02e4b0 100644 --- a/rules/parser.go +++ b/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":