hayzam
7a3f6e945d
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.
91 lines
2.4 KiB
Go
91 lines
2.4 KiB
Go
package ruleset
|
|
|
|
import (
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
|
|
"git.difuse.io/Difuse/Mellaris/analyzer"
|
|
"git.difuse.io/Difuse/Mellaris/ruleset/builtins/geo"
|
|
|
|
"github.com/expr-lang/expr/ast"
|
|
"github.com/expr-lang/expr/parser"
|
|
)
|
|
|
|
func TestExtractGeoSiteConditions(t *testing.T) {
|
|
expression := `
|
|
(geosite(tls.req.sni, "openai") || geosite(quic.req.sni, "OpenAI")) &&
|
|
geosite(http.req.headers.host, "google@ads")
|
|
`
|
|
got := extractGeoSiteConditions(expression)
|
|
want := []string{"openai", "google@ads"}
|
|
if !reflect.DeepEqual(got, want) {
|
|
t.Fatalf("extractGeoSiteConditions() = %v, want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestExtractGeoSiteHostCandidates(t *testing.T) {
|
|
info := StreamInfo{
|
|
Props: analyzer.CombinedPropMap{
|
|
"quic": analyzer.PropMap{
|
|
"req": analyzer.PropMap{
|
|
"sni": "ChatGPT.com",
|
|
},
|
|
},
|
|
"http": analyzer.PropMap{
|
|
"req": analyzer.PropMap{
|
|
"headers": analyzer.PropMap{
|
|
"host": "api.openai.com:443",
|
|
},
|
|
},
|
|
},
|
|
"dns": analyzer.PropMap{
|
|
"questions": []analyzer.PropMap{
|
|
{"name": "chatgpt.com."},
|
|
{"name": "8.8.8.8"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
got := extractGeoSiteHostCandidates(info)
|
|
want := []string{"chatgpt.com", "api.openai.com"}
|
|
if !reflect.DeepEqual(got, want) {
|
|
t.Fatalf("extractGeoSiteHostCandidates() = %v, want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestMatchGeoSiteConditions(t *testing.T) {
|
|
hosts := []string{"chatgpt.com", "api.openai.com"}
|
|
conditions := []string{" openai ", "google", "OPENAI"}
|
|
got := matchGeoSiteConditions(hosts, conditions, func(site, condition string) bool {
|
|
if condition != "openai" {
|
|
return false
|
|
}
|
|
return site == "chatgpt.com" || site == "api.openai.com"
|
|
})
|
|
want := []string{"openai"}
|
|
if !reflect.DeepEqual(got, want) {
|
|
t.Fatalf("matchGeoSiteConditions() = %v, want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestIDPatcher_PatchesGeoSiteORChainToGeoSiteSet(t *testing.T) {
|
|
tree, err := parser.Parse(`geosite(tls.req.sni, "google") || geosite(tls.req.sni, "youtube") || geosite(tls.req.sni, "openai")`)
|
|
if err != nil {
|
|
t.Fatalf("parse expression: %v", err)
|
|
}
|
|
root := tree.Node
|
|
patcher := &idPatcher{GeoMatcher: geo.NewGeoMatcher("", "")}
|
|
ast.Walk(&root, patcher)
|
|
if patcher.Err != nil {
|
|
t.Fatalf("patch error: %v", patcher.Err)
|
|
}
|
|
got := root.String()
|
|
if !strings.Contains(got, "geosite_set(") {
|
|
t.Fatalf("expected geosite_set rewrite, got %q", got)
|
|
}
|
|
if strings.Contains(got, "||") || strings.Contains(got, " or ") {
|
|
t.Fatalf("expected OR chain to be collapsed, got %q", got)
|
|
}
|
|
}
|