chore: Add source matching for ip type rules

MetaCubeX · Apr 1, 2024 · 3b472f7 · 3b472f7
1 parent f3743fc
commit 3b472f7
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 12 deletions.
diff --git 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
@@ -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,
  }

diff --git 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
@@ -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":