Improves flow handling and adds runtime stats APIs
Refactors TCP and UDP flow managers to enhance analyzer selection and flow binding accuracy, including O(1) UDP stream rebinding by 5-tuple. Introduces runtime stats tracking for engine and ruleset operations, exposing new APIs for granular performance and error metrics. Optimizes GeoMatcher with result caching and supports efficient geosite set matching, reducing redundant computation in ruleset expressions.
This commit is contained in:
@@ -1,52 +1,45 @@
|
||||
package geo
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"net"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
const (
|
||||
geoSiteResultCacheSize = 1 << 16
|
||||
geoSiteSetResultCacheSize = 1 << 16
|
||||
)
|
||||
|
||||
type GeoMatcher struct {
|
||||
geoLoader GeoLoader
|
||||
geoSiteMatcher map[string]hostMatcher
|
||||
siteMatcherLock sync.Mutex
|
||||
siteMatcherLock sync.RWMutex
|
||||
geoSiteSets map[string][]hostMatcher
|
||||
siteSetLock sync.RWMutex
|
||||
geoIpMatcher map[string]hostMatcher
|
||||
ipMatcherLock sync.Mutex
|
||||
ipMatcherLock sync.RWMutex
|
||||
geoSiteResult *boolLRUCache
|
||||
geoSiteSetCache *boolLRUCache
|
||||
}
|
||||
|
||||
func NewGeoMatcher(geoSiteFilename, geoIpFilename string) *GeoMatcher {
|
||||
return &GeoMatcher{
|
||||
geoLoader: NewDefaultGeoLoader(geoSiteFilename, geoIpFilename),
|
||||
geoSiteMatcher: make(map[string]hostMatcher),
|
||||
geoIpMatcher: make(map[string]hostMatcher),
|
||||
geoLoader: NewDefaultGeoLoader(geoSiteFilename, geoIpFilename),
|
||||
geoSiteMatcher: make(map[string]hostMatcher),
|
||||
geoSiteSets: make(map[string][]hostMatcher),
|
||||
geoIpMatcher: make(map[string]hostMatcher),
|
||||
geoSiteResult: newBoolLRUCache(geoSiteResultCacheSize),
|
||||
geoSiteSetCache: newBoolLRUCache(geoSiteSetResultCacheSize),
|
||||
}
|
||||
}
|
||||
|
||||
func (g *GeoMatcher) MatchGeoIp(ip, condition string) bool {
|
||||
g.ipMatcherLock.Lock()
|
||||
defer g.ipMatcherLock.Unlock()
|
||||
|
||||
matcher, ok := g.geoIpMatcher[condition]
|
||||
if !ok {
|
||||
// GeoIP matcher
|
||||
condition = strings.ToLower(condition)
|
||||
country := condition
|
||||
if len(country) == 0 {
|
||||
return false
|
||||
}
|
||||
gMap, err := g.geoLoader.LoadGeoIP()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
list, ok := gMap[country]
|
||||
if !ok || list == nil {
|
||||
return false
|
||||
}
|
||||
matcher, err = newGeoIPMatcher(list)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
g.geoIpMatcher[condition] = matcher
|
||||
matcher, ok := g.getOrCreateGeoIPMatcher(condition)
|
||||
if !ok || matcher == nil {
|
||||
return false
|
||||
}
|
||||
parseIp := net.ParseIP(ip)
|
||||
if parseIp == nil {
|
||||
@@ -64,32 +57,69 @@ func (g *GeoMatcher) MatchGeoIp(ip, condition string) bool {
|
||||
}
|
||||
|
||||
func (g *GeoMatcher) MatchGeoSite(site, condition string) bool {
|
||||
g.siteMatcherLock.Lock()
|
||||
defer g.siteMatcherLock.Unlock()
|
||||
|
||||
matcher, ok := g.geoSiteMatcher[condition]
|
||||
if !ok {
|
||||
// MatchGeoSite matcher
|
||||
condition = strings.ToLower(condition)
|
||||
name, attrs := parseGeoSiteName(condition)
|
||||
if len(name) == 0 {
|
||||
return false
|
||||
}
|
||||
gMap, err := g.geoLoader.LoadGeoSite()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
list, ok := gMap[name]
|
||||
if !ok || list == nil {
|
||||
return false
|
||||
}
|
||||
matcher, err = newGeositeMatcher(list, attrs)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
g.geoSiteMatcher[condition] = matcher
|
||||
conditionKey := strings.TrimSpace(strings.ToLower(condition))
|
||||
if conditionKey == "" {
|
||||
return false
|
||||
}
|
||||
return matcher.Match(HostInfo{Name: site})
|
||||
cacheKey := site + "\x1f" + conditionKey
|
||||
if v, ok := g.geoSiteResult.Get(cacheKey); ok {
|
||||
return v
|
||||
}
|
||||
matcher, ok := g.getOrCreateGeoSiteMatcher(condition)
|
||||
if !ok || matcher == nil {
|
||||
return false
|
||||
}
|
||||
result := matcher.Match(HostInfo{Name: site})
|
||||
g.geoSiteResult.Set(cacheKey, result)
|
||||
return result
|
||||
}
|
||||
|
||||
func (g *GeoMatcher) MatchGeoSiteSet(site string, set *SiteConditionSet) bool {
|
||||
if set == nil {
|
||||
return false
|
||||
}
|
||||
conditions := normalizeGeoSiteSetConditions(set.Conditions)
|
||||
if len(conditions) == 0 {
|
||||
return false
|
||||
}
|
||||
key := strings.Join(conditions, "\x1f")
|
||||
cacheKey := site + "\x1e" + key
|
||||
if v, ok := g.geoSiteSetCache.Get(cacheKey); ok {
|
||||
return v
|
||||
}
|
||||
|
||||
g.siteSetLock.RLock()
|
||||
matchers, ok := g.geoSiteSets[key]
|
||||
g.siteSetLock.RUnlock()
|
||||
if !ok {
|
||||
compiled := make([]hostMatcher, 0, len(conditions))
|
||||
for _, condition := range conditions {
|
||||
m, ok := g.getOrCreateGeoSiteMatcher(condition)
|
||||
if ok && m != nil {
|
||||
compiled = append(compiled, m)
|
||||
}
|
||||
}
|
||||
g.siteSetLock.Lock()
|
||||
if existing, exists := g.geoSiteSets[key]; exists {
|
||||
matchers = existing
|
||||
} else {
|
||||
g.geoSiteSets[key] = compiled
|
||||
matchers = compiled
|
||||
}
|
||||
g.siteSetLock.Unlock()
|
||||
}
|
||||
if len(matchers) == 0 {
|
||||
return false
|
||||
}
|
||||
host := HostInfo{Name: site}
|
||||
for _, matcher := range matchers {
|
||||
if matcher.Match(host) {
|
||||
g.geoSiteSetCache.Set(cacheKey, true)
|
||||
return true
|
||||
}
|
||||
}
|
||||
g.geoSiteSetCache.Set(cacheKey, false)
|
||||
return false
|
||||
}
|
||||
|
||||
func (g *GeoMatcher) LoadGeoSite() error {
|
||||
@@ -111,3 +141,152 @@ func parseGeoSiteName(s string) (string, []string) {
|
||||
}
|
||||
return base, attrs
|
||||
}
|
||||
|
||||
func (g *GeoMatcher) getOrCreateGeoSiteMatcher(condition string) (hostMatcher, bool) {
|
||||
condition = strings.TrimSpace(strings.ToLower(condition))
|
||||
if condition == "" {
|
||||
return nil, false
|
||||
}
|
||||
g.siteMatcherLock.RLock()
|
||||
matcher, ok := g.geoSiteMatcher[condition]
|
||||
g.siteMatcherLock.RUnlock()
|
||||
if ok {
|
||||
return matcher, true
|
||||
}
|
||||
|
||||
name, attrs := parseGeoSiteName(condition)
|
||||
if len(name) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
gMap, err := g.geoLoader.LoadGeoSite()
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
list, ok := gMap[name]
|
||||
if !ok || list == nil {
|
||||
return nil, false
|
||||
}
|
||||
matcher, err = newGeositeMatcher(list, attrs)
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
g.siteMatcherLock.Lock()
|
||||
if existing, exists := g.geoSiteMatcher[condition]; exists {
|
||||
matcher = existing
|
||||
} else {
|
||||
g.geoSiteMatcher[condition] = matcher
|
||||
}
|
||||
g.siteMatcherLock.Unlock()
|
||||
return matcher, true
|
||||
}
|
||||
|
||||
func (g *GeoMatcher) getOrCreateGeoIPMatcher(condition string) (hostMatcher, bool) {
|
||||
condition = strings.TrimSpace(strings.ToLower(condition))
|
||||
if condition == "" {
|
||||
return nil, false
|
||||
}
|
||||
g.ipMatcherLock.RLock()
|
||||
matcher, ok := g.geoIpMatcher[condition]
|
||||
g.ipMatcherLock.RUnlock()
|
||||
if ok {
|
||||
return matcher, true
|
||||
}
|
||||
|
||||
gMap, err := g.geoLoader.LoadGeoIP()
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
list, ok := gMap[condition]
|
||||
if !ok || list == nil {
|
||||
return nil, false
|
||||
}
|
||||
matcher, err = newGeoIPMatcher(list)
|
||||
if err != nil {
|
||||
return nil, false
|
||||
}
|
||||
g.ipMatcherLock.Lock()
|
||||
if existing, exists := g.geoIpMatcher[condition]; exists {
|
||||
matcher = existing
|
||||
} else {
|
||||
g.geoIpMatcher[condition] = matcher
|
||||
}
|
||||
g.ipMatcherLock.Unlock()
|
||||
return matcher, true
|
||||
}
|
||||
|
||||
func normalizeGeoSiteSetConditions(in []string) []string {
|
||||
if len(in) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(in))
|
||||
seen := make(map[string]struct{}, len(in))
|
||||
for _, v := range in {
|
||||
s := strings.TrimSpace(strings.ToLower(v))
|
||||
if s == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[s]; ok {
|
||||
continue
|
||||
}
|
||||
seen[s] = struct{}{}
|
||||
out = append(out, s)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
type boolLRUCache struct {
|
||||
mu sync.Mutex
|
||||
cap int
|
||||
ll *list.List
|
||||
items map[string]*list.Element
|
||||
}
|
||||
|
||||
type boolCacheEntry struct {
|
||||
key string
|
||||
value bool
|
||||
}
|
||||
|
||||
func newBoolLRUCache(capacity int) *boolLRUCache {
|
||||
if capacity <= 0 {
|
||||
capacity = 1
|
||||
}
|
||||
return &boolLRUCache{
|
||||
cap: capacity,
|
||||
ll: list.New(),
|
||||
items: make(map[string]*list.Element, capacity),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *boolLRUCache) Get(key string) (bool, bool) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if ele, ok := c.items[key]; ok {
|
||||
c.ll.MoveToFront(ele)
|
||||
entry := ele.Value.(boolCacheEntry)
|
||||
return entry.value, true
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
|
||||
func (c *boolLRUCache) Set(key string, value bool) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if ele, ok := c.items[key]; ok {
|
||||
ele.Value = boolCacheEntry{key: key, value: value}
|
||||
c.ll.MoveToFront(ele)
|
||||
return
|
||||
}
|
||||
ele := c.ll.PushFront(boolCacheEntry{key: key, value: value})
|
||||
c.items[key] = ele
|
||||
if c.ll.Len() <= c.cap {
|
||||
return
|
||||
}
|
||||
back := c.ll.Back()
|
||||
if back == nil {
|
||||
return
|
||||
}
|
||||
entry := back.Value.(boolCacheEntry)
|
||||
delete(c.items, entry.key)
|
||||
c.ll.Remove(back)
|
||||
}
|
||||
|
||||
@@ -1,13 +1,14 @@
|
||||
package geo
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"git.difuse.io/Difuse/Mellaris/ruleset/builtins/geo/v2geo"
|
||||
)
|
||||
|
||||
type fakeGeoLoader struct {
|
||||
geoip map[string]*v2geo.GeoIP
|
||||
geoip map[string]*v2geo.GeoIP
|
||||
geosite map[string]*v2geo.GeoSite
|
||||
}
|
||||
|
||||
@@ -110,6 +111,83 @@ func TestGeoMatcher_MatchGeoSite_MissingSite(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeoMatcher_MatchGeoSiteSet(t *testing.T) {
|
||||
loader := &fakeGeoLoader{
|
||||
geosite: map[string]*v2geo.GeoSite{
|
||||
"openai": {
|
||||
Domain: []*v2geo.Domain{
|
||||
{Type: v2geo.Domain_Plain, Value: "openai"},
|
||||
},
|
||||
},
|
||||
"google": {
|
||||
Domain: []*v2geo.Domain{
|
||||
{Type: v2geo.Domain_RootDomain, Value: "google.com"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
g := NewGeoMatcher("", "")
|
||||
g.geoLoader = loader
|
||||
|
||||
set := &SiteConditionSet{Conditions: []string{" google ", "openai", "OPENAI"}}
|
||||
if !g.MatchGeoSiteSet("api.openai.com", set) {
|
||||
t.Error("MatchGeoSiteSet should match openai")
|
||||
}
|
||||
if !g.MatchGeoSiteSet("mail.google.com", set) {
|
||||
t.Error("MatchGeoSiteSet should match google")
|
||||
}
|
||||
if g.MatchGeoSiteSet("example.com", set) {
|
||||
t.Error("MatchGeoSiteSet should not match unrelated host")
|
||||
}
|
||||
}
|
||||
|
||||
type countingMatcher struct {
|
||||
calls *atomic.Uint64
|
||||
match bool
|
||||
}
|
||||
|
||||
func (m countingMatcher) Match(host HostInfo) bool {
|
||||
_ = host
|
||||
m.calls.Add(1)
|
||||
return m.match
|
||||
}
|
||||
|
||||
func TestGeoMatcher_MatchGeoSite_UsesResultCache(t *testing.T) {
|
||||
g := NewGeoMatcher("", "")
|
||||
var calls atomic.Uint64
|
||||
g.geoSiteMatcher["openai"] = countingMatcher{calls: &calls, match: true}
|
||||
|
||||
if !g.MatchGeoSite("api.openai.com", "openai") {
|
||||
t.Fatal("expected match")
|
||||
}
|
||||
if !g.MatchGeoSite("api.openai.com", "openai") {
|
||||
t.Fatal("expected cached match")
|
||||
}
|
||||
if got := calls.Load(); got != 1 {
|
||||
t.Fatalf("matcher calls=%d want=1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeoMatcher_MatchGeoSiteSet_UsesResultCache(t *testing.T) {
|
||||
g := NewGeoMatcher("", "")
|
||||
var calls atomic.Uint64
|
||||
g.geoSiteSets["openai\x1fyoutube"] = []hostMatcher{
|
||||
countingMatcher{calls: &calls, match: false},
|
||||
countingMatcher{calls: &calls, match: true},
|
||||
}
|
||||
|
||||
set := &SiteConditionSet{Conditions: []string{"youtube", "openai"}}
|
||||
if !g.MatchGeoSiteSet("www.youtube.com", set) {
|
||||
t.Fatal("expected match")
|
||||
}
|
||||
if !g.MatchGeoSiteSet("www.youtube.com", set) {
|
||||
t.Fatal("expected cached match")
|
||||
}
|
||||
if got := calls.Load(); got != 2 {
|
||||
t.Fatalf("matcher calls=%d want=2", got)
|
||||
}
|
||||
}
|
||||
|
||||
func ipv4(a, b, c, d byte) []byte {
|
||||
return []byte{a, b, c, d}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,10 @@ type HostInfo struct {
|
||||
IPv6 net.IP
|
||||
}
|
||||
|
||||
type SiteConditionSet struct {
|
||||
Conditions []string
|
||||
}
|
||||
|
||||
func (h HostInfo) String() string {
|
||||
return fmt.Sprintf("%s|%s|%s", h.Name, h.IPv4, h.IPv6)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user