MigrateRepos/OpenGFW
1
0
Fork 0
mirror of https://github.com/apernet/OpenGFW.git synced 2024-12-22 08:59:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: dns lookup function

Browse Source
This commit is contained in:
Toby 2024-04-03 20:02:57 -07:00
parent d7737e9211
commit ae34b4856a
4 changed files with 115 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
go.mod
View File

@ -5,7 +5,7 @@ go 1.21
require (
github.com/bwmarrin/snowflake v0.3.0
github.com/coreos/go-iptables v0.7.0
github.com/expr-lang/expr v1.15.7
github.com/expr-lang/expr v1.16.3
github.com/florianl/go-nfqueue v1.3.2-0.20231218173729-f2bdeb033acf
github.com/google/gopacket v1.1.20-0.20220810144506-32ee38206866
github.com/hashicorp/golang-lru/v2 v2.0.7

4
go.sum
View File

@ -7,8 +7,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/expr-lang/expr v1.15.7 h1:BK0JcWUkoW6nrbLBo6xCKhz4BvH5DSOOu1Gx5lucyZo=
github.com/expr-lang/expr v1.15.7/go.mod h1:uCkhfG+x7fcZ5A5sXHKuQ07jGZRl6J0FCAaf2k4PtVQ=
github.com/expr-lang/expr v1.16.3 h1:NLldf786GffptcXNxxJx5dQ+FzeWDKChBDqOOwyK8to=
github.com/expr-lang/expr v1.16.3/go.mod h1:uCkhfG+x7fcZ5A5sXHKuQ07jGZRl6J0FCAaf2k4PtVQ=
github.com/florianl/go-nfqueue v1.3.2-0.20231218173729-f2bdeb033acf h1:NqGS3vTHzVENbIfd87cXZwdpO6MB2R1PjHMJLi4Z3ow=
github.com/florianl/go-nfqueue v1.3.2-0.20231218173729-f2bdeb033acf/go.mod h1:eSnAor2YCfMCVYrVNEhkLGN/r1L+J4uDjc0EUy0tfq4=
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=

2
io/nfqueue.go
View File

@ -55,7 +55,7 @@ func generateNftRules(local, rst bool) (*nftTableSpec, error) {
c.Rules = append(c.Rules, "ip protocol tcp ct mark $DROP_CTMARK counter reject with tcp reset")
}
c.Rules = append(c.Rules, "ct mark $DROP_CTMARK counter drop")
c.Rules = append(c.Rules, "counter queue num $QUEUE_NUM bypass")
c.Rules = append(c.Rules, "ip protocol tcp counter queue num $QUEUE_NUM bypass")
}
return table, nil
}

159
ruleset/expr.go
View File

@ -1,11 +1,15 @@
package ruleset
import (
"context"
"fmt"
"net"
"os"
"reflect"
"strings"
"time"
"github.com/expr-lang/expr/builtin"
"github.com/expr-lang/expr"
"github.com/expr-lang/expr/ast"
@ -104,6 +108,7 @@ func CompileExprRules(rules []ExprRule, ans []analyzer.Analyzer, mods []modifier
if err != nil {
return nil, err
}
funcMap := buildFunctionMap(geoMatcher)
// Compile all rules and build a map of analyzers that are used by the rules.
for _, rule := range rules {
if rule.Action == "" && !rule.Log {
@ -118,13 +123,19 @@ func CompileExprRules(rules []ExprRule, ans []analyzer.Analyzer, mods []modifier
action = &a
}
visitor := &idVisitor{Variables: make(map[string]bool), Identifiers: make(map[string]bool)}
patcher := &idPatcher{}
patcher := &idPatcher{FuncMap: funcMap}
program, err := expr.Compile(rule.Expr,
func(c *conf.Config) {
c.Strict = false
c.Expect = reflect.Bool
c.Visitors = append(c.Visitors, visitor, patcher)
registerBuiltinFunctions(c.Functions, geoMatcher)
for name, f := range funcMap {
c.Functions[name] = &builtin.Function{
Name: name,
Func: f.Func,
Types: f.Types,
}
}
},
)
if err != nil {
@ -138,24 +149,15 @@ func CompileExprRules(rules []ExprRule, ans []analyzer.Analyzer, mods []modifier
if isBuiltInAnalyzer(name) || visitor.Variables[name] {
continue
}
// Check if it's one of the built-in functions, and if so,
// skip it as an analyzer & do initialization if necessary.
switch name {
case "geoip":
if err := geoMatcher.LoadGeoIP(); err != nil {
return nil, fmt.Errorf("rule %q failed to load geoip: %w", rule.Name, err)
if f, ok := funcMap[name]; ok {
// Built-in function, initialize if necessary
if f.InitFunc != nil {
if err := f.InitFunc(); err != nil {
return nil, fmt.Errorf("rule %q failed to initialize function %q: %w", rule.Name, name, err)
}
case "geosite":
if err := geoMatcher.LoadGeoSite(); err != nil {
return nil, fmt.Errorf("rule %q failed to load geosite: %w", rule.Name, err)
}
case "cidr":
// No initialization needed for CIDR.
default:
a, ok := fullAnMap[name]
if !ok {
return nil, fmt.Errorf("rule %q uses unknown analyzer %q", rule.Name, name)
}
} else if a, ok := fullAnMap[name]; ok {
// Analyzer, add to dependency map
depAnMap[name] = a
}
}
@ -191,30 +193,6 @@ func CompileExprRules(rules []ExprRule, ans []analyzer.Analyzer, mods []modifier
}, nil
}
func registerBuiltinFunctions(funcMap map[string]*ast.Function, geoMatcher *geo.GeoMatcher) {
funcMap["geoip"] = &ast.Function{
Name: "geoip",
Func: func(params ...any) (any, error) {
return geoMatcher.MatchGeoIp(params[0].(string), params[1].(string)), nil
},
Types: []reflect.Type{reflect.TypeOf(geoMatcher.MatchGeoIp)},
}
funcMap["geosite"] = &ast.Function{
Name: "geosite",
Func: func(params ...any) (any, error) {
return geoMatcher.MatchGeoSite(params[0].(string), params[1].(string)), nil
},
Types: []reflect.Type{reflect.TypeOf(geoMatcher.MatchGeoSite)},
}
funcMap["cidr"] = &ast.Function{
Name: "cidr",
Func: func(params ...any) (any, error) {
return builtins.MatchCIDR(params[0].(string), params[1].(*net.IPNet)), nil
},
Types: []reflect.Type{reflect.TypeOf((func(string, string) bool)(nil)), reflect.TypeOf(builtins.MatchCIDR)},
}
}
func streamInfoToExprEnv(info StreamInfo) map[string]interface{} {
m := map[string]interface{}{
"id": info.ID,
@ -299,6 +277,7 @@ func (v *idVisitor) Visit(node *ast.Node) {
// idPatcher patches the AST during expr compilation, replacing certain values with
// their internal representations for better runtime performance.
type idPatcher struct {
FuncMap map[string]*Function
Err error
}
@ -306,22 +285,98 @@ func (p *idPatcher) Visit(node *ast.Node) {
switch (*node).(type) {
case *ast.CallNode:
callNode := (*node).(*ast.CallNode)
if callNode.Func == nil {
if callNode.Callee == nil {
// Ignore invalid call nodes
return
}
switch callNode.Func.Name {
case "cidr":
cidrStringNode, ok := callNode.Arguments[1].(*ast.StringNode)
if !ok {
return
}
cidr, err := builtins.CompileCIDR(cidrStringNode.Value)
if err != nil {
if f, ok := p.FuncMap[callNode.Callee.String()]; ok {
if f.PatchFunc != nil {
if err := f.PatchFunc(&callNode.Arguments); err != nil {
p.Err = err
return
}
callNode.Arguments[1] = &ast.ConstantNode{Value: cidr}
}
}
}
}
type Function struct {
InitFunc func() error
PatchFunc func(args *[]ast.Node) error
Func func(params ...any) (any, error)
Types []reflect.Type
}
func buildFunctionMap(geoMatcher *geo.GeoMatcher) map[string]*Function {
return map[string]*Function{
"geoip": {
InitFunc: geoMatcher.LoadGeoIP,
PatchFunc: nil,
Func: func(params ...any) (any, error) {
return geoMatcher.MatchGeoIp(params[0].(string), params[1].(string)), nil
},
Types: []reflect.Type{reflect.TypeOf(geoMatcher.MatchGeoIp)},
},
"geosite": {
InitFunc: geoMatcher.LoadGeoSite,
PatchFunc: nil,
Func: func(params ...any) (any, error) {
return geoMatcher.MatchGeoSite(params[0].(string), params[1].(string)), nil
},
Types: []reflect.Type{reflect.TypeOf(geoMatcher.MatchGeoSite)},
},
"cidr": {
InitFunc: nil,
PatchFunc: func(args *[]ast.Node) error {
cidrStringNode, ok := (*args)[1].(*ast.StringNode)
if !ok {
return fmt.Errorf("cidr: invalid argument type")
}
cidr, err := builtins.CompileCIDR(cidrStringNode.Value)
if err != nil {
return err
}
(*args)[1] = &ast.ConstantNode{Value: cidr}
return nil
},
Func: func(params ...any) (any, error) {
return builtins.MatchCIDR(params[0].(string), params[1].(*net.IPNet)), nil
},
Types: []reflect.Type{reflect.TypeOf((func(string, string) bool)(nil)), reflect.TypeOf(builtins.MatchCIDR)},
},
"lookup": {
InitFunc: nil,
PatchFunc: func(args *[]ast.Node) error {
if len(*args) < 2 {
// Second argument (DNS server) is optional
return nil
}
serverStr, ok := (*args)[1].(*ast.StringNode)
if !ok {
return fmt.Errorf("lookup: invalid argument type")
}
r := &net.Resolver{
Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
return net.Dial(network, serverStr.Value)
},
}
(*args)[1] = &ast.ConstantNode{Value: r}
return nil
},
Func: func(params ...any) (any, error) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if len(params) < 2 {
return net.DefaultResolver.LookupHost(ctx, params[0].(string))
} else {
return params[1].(*net.Resolver).LookupHost(ctx, params[0].(string))
}
},
Types: []reflect.Type{
reflect.TypeOf((func(string, string) []string)(nil)),
reflect.TypeOf((func(string) []string)(nil)),
reflect.TypeOf((func(string, *net.Resolver) []string)(nil)),
},
},
}
}