feat: dns lookup function

Author: tobyxdd

go.mod

@@ -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

go.sum

@@ -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=

io/nfqueue.go

@@ -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
 }

ruleset/expr.go

@@ -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)
-				}
-			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)
+			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)
+					}
 				}
+			} 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,29 +277,106 @@ 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 {
-	Err error
+	FuncMap map[string]*Function
+	Err     error
 }
 
 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 {
-				p.Err = err
-				return
+		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)),
+			},
+		},
+	}
+}