feat(dns): add Lua scripting for DNS queries

app/dns/script.go

@@ -0,0 +1,71 @@
+package dns
+
+import (
+	"context"
+	"time"
+
+	"github.com/xtls/xray-core/common/errors"
+	"github.com/xtls/xray-core/common/geodata"
+	luamgr "github.com/xtls/xray-core/common/lua"
+	"github.com/xtls/xray-core/common/net"
+	"github.com/xtls/xray-core/features/dns"
+	lua "github.com/yuin/gopher-lua"
+)
+
+const scriptExecutionTimeout = 10 * time.Second
+
+type scriptEngine struct {
+	dns  *DNS
+	pool *luamgr.Pool
+}
+
+func newScriptEngine(path string, server *DNS) (*scriptEngine, error) {
+	program, err := luamgr.CompileFile(path)
+	if err != nil {
+		return nil, err
+	}
+	e := &scriptEngine{dns: server}
+	e.pool, err = luamgr.NewPool(server.ctx, func(poolCtx context.Context) (*lua.LState, error) {
+		initCtx, cancel := context.WithTimeout(poolCtx, scriptExecutionTimeout)
+		defer cancel()
+		L, err := program.NewState(initCtx, func(L *lua.LState) {
+			geodata.RegisterLua(L)
+			server.RegisterLua(L)
+		})
+		if err != nil {
+			return nil, err
+		}
+		if L.GetGlobal("handleDNSQuery").Type() != lua.LTFunction {
+			L.Close()
+			return nil, errors.New("DNS script must define handleDNSQuery(q)")
+		}
+		return L, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	errors.LogInfo(server.ctx, "DNS script initialized from ", path)
+	return e, nil
+}
+
+func (e *scriptEngine) close() {
+	e.pool.Close()
+}
+
+func (e *scriptEngine) query(domain string, option dns.IPOption) ([]net.IP, uint32, error) {
+	L, err := e.pool.Acquire()
+	if err != nil {
+		return nil, 0, err
+	}
+	reusable := false
+	defer func() {
+		e.pool.Release(L, reusable)
+	}()
+	queryCtx, cancel := context.WithTimeout(e.pool.Context(), scriptExecutionTimeout)
+	defer cancel()
+	ips, ttl, err := e.dns.CallLuaHook(L, queryCtx, domain, option)
+	if err == nil {
+		reusable = true
+	}
+	return ips, ttl, err
+}