feat(dns): add Lua scripting for DNS queries

common/geodata/lua.go

@@ -0,0 +1,60 @@
+package geodata
+
+import (
+	lua "github.com/yuin/gopher-lua"
+	luar "layeh.com/gopher-luar"
+)
+
+// RegisterLua makes xray.geodata available to require in an LState.
+// Matchers retain registry handles, so they remain usable after a reload.
+func RegisterLua(L *lua.LState) {
+	L.PreloadModule("xray.geodata", func(L *lua.LState) int {
+		module := L.NewTable()
+
+		module.RawSetString("domainMatcher", L.NewFunction(func(L *lua.LState) int {
+			parsed, err := ParseDomainRules(luaRules(L, 1), Domain_Domain)
+			if err != nil {
+				L.RaiseError("%v", err)
+				return 0
+			}
+			matcher, err := DomainReg.BuildDomainMatcher(parsed)
+			if err != nil {
+				L.RaiseError("%v", err)
+				return 0
+			}
+			L.Push(luar.New(L, matcher))
+			return 1
+		}))
+
+		module.RawSetString("ipMatcher", L.NewFunction(func(L *lua.LState) int {
+			parsed, err := ParseIPRules(luaRules(L, 1))
+			if err != nil {
+				L.RaiseError("%v", err)
+				return 0
+			}
+			matcher, err := IPReg.BuildIPMatcher(parsed)
+			if err != nil {
+				L.RaiseError("%v", err)
+				return 0
+			}
+			L.Push(luar.New(L, matcher))
+			return 1
+		}))
+		L.Push(module)
+		return 1
+	})
+}
+
+func luaRules(L *lua.LState, index int) []string {
+	table := L.CheckTable(index)
+	rules := make([]string, table.Len())
+	for i := range rules {
+		value, ok := table.RawGetInt(i + 1).(lua.LString)
+		if !ok {
+			L.RaiseError("geodata rules must be strings")
+			return nil
+		}
+		rules[i] = string(value)
+	}
+	return rules
+}

common/geodata/lua_test.go

@@ -0,0 +1,42 @@
+package geodata
+
+import (
+	"testing"
+
+	"github.com/xtls/xray-core/common/net"
+	lua "github.com/yuin/gopher-lua"
+)
+
+func TestLuaIPMatcherAcceptsNativeIP(t *testing.T) {
+	L := lua.NewState()
+	defer L.Close()
+	RegisterLua(L)
+	ip := L.NewUserData()
+	ip.Value = net.ParseIP("127.0.0.1")
+	L.SetGlobal("ip", ip)
+	if err := L.DoString(`
+		local matcher = require("xray.geodata").ipMatcher({"127.0.0.0/8"})
+		assert(matcher:Match(ip))
+		assert(matcher:AnyMatch({ip}))
+		assert(matcher:Matches({ip}))
+		local matched, unmatched = matcher:FilterIPs({ip})
+		assert(#matched == 1 and #unmatched == 0)
+		assert(matcher:Match(matched[1]))
+	`); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestLuaDomainMatcherUsesNativeMatcher(t *testing.T) {
+	L := lua.NewState()
+	defer L.Close()
+	RegisterLua(L)
+	if err := L.DoString(`
+		local matcher = require("xray.geodata").domainMatcher({"example.com"})
+		assert(matcher:MatchAny("example.com"))
+		assert(matcher:MatchAny("www.example.com"))
+		assert(#(matcher:Match("www.example.com")) == 1)
+	`); err != nil {
+		t.Fatal(err)
+	}
+}

common/lua/pool.go

@@ -0,0 +1,138 @@
+package lua
+
+import (
+	"context"
+	"errors"
+	"sync"
+
+	glua "github.com/yuin/gopher-lua"
+)
+
+const maxIdleStates = 16
+
+// LStateFactory must initialize a state fully and observe ctx while doing so.
+// The pool owns any non-nil state it returns, even when it also returns an error.
+type LStateFactory func(ctx context.Context) (*glua.LState, error)
+
+// Pool lends each state to one caller at a time. It grows on contention and
+// keeps up to maxIdleStates idle states until Close. Callers decide whether a
+// state is reusable.
+type Pool struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+
+	factory LStateFactory
+	idle    []*glua.LState
+
+	mu     sync.Mutex
+	active sync.WaitGroup
+	closed bool
+}
+
+// NewPool initializes one state before returning, so top-level errors surface at startup.
+func NewPool(ctx context.Context, factory LStateFactory) (*Pool, error) {
+	poolCtx, cancel := context.WithCancel(ctx)
+
+	// Create one state now to catch factory errors at startup.
+	state, err := factory(poolCtx)
+	if err != nil {
+		cancel()
+		if state != nil {
+			state.Close()
+		}
+		return nil, err
+	}
+	if state == nil {
+		cancel()
+		return nil, errors.New("Lua state factory returned nil")
+	}
+	if err := poolCtx.Err(); err != nil {
+		state.Close()
+		cancel()
+		return nil, err
+	}
+
+	return &Pool{ctx: poolCtx, cancel: cancel, factory: factory, idle: []*glua.LState{state}}, nil
+}
+
+// Context is cancelled by Close. Query contexts should derive from it.
+func (p *Pool) Context() context.Context {
+	return p.ctx
+}
+
+// Acquire returns an initialized exclusive state, growing the pool if necessary.
+func (p *Pool) Acquire() (*glua.LState, error) {
+	p.mu.Lock()
+	if p.closed || p.ctx.Err() != nil {
+		p.mu.Unlock()
+		return nil, p.ctx.Err()
+	}
+
+	p.active.Add(1)
+
+	n := len(p.idle)
+	if n != 0 {
+		state := p.idle[n-1]
+		p.idle = p.idle[:n-1]
+		p.mu.Unlock()
+		return state, nil
+	}
+	p.mu.Unlock()
+
+	// TODO: Limit the total number of states. When the limit is reached, wait
+	// for a Release instead of creating another state; allow the wait to be
+	// cancelled by the caller or by Close.
+	state, err := p.factory(p.ctx)
+	if err == nil && state == nil {
+		err = errors.New("Lua state factory returned nil")
+	}
+	if err != nil {
+		if state != nil {
+			state.Close()
+		}
+		p.active.Done()
+		return nil, err
+	}
+	if err := p.ctx.Err(); err != nil {
+		state.Close()
+		p.active.Done()
+		return nil, err
+	}
+
+	return state, nil
+}
+
+// Release returns a healthy state to the pool and closes a failed or cancelled one.
+func (p *Pool) Release(state *glua.LState, reusable bool) {
+	if reusable {
+		p.mu.Lock()
+		if !p.closed && p.ctx.Err() == nil && len(p.idle) < maxIdleStates {
+			p.idle = append(p.idle, state)
+		} else {
+			reusable = false
+		}
+		p.mu.Unlock()
+	}
+
+	if !reusable {
+		state.Close()
+	}
+
+	p.active.Done()
+}
+
+// Close cancels active work, closes idle states, and waits for borrowed states.
+func (p *Pool) Close() {
+	p.mu.Lock()
+	if !p.closed {
+		p.closed = true
+		p.cancel()
+		for _, state := range p.idle {
+			state.Close()
+		}
+		p.idle = nil
+	}
+	p.mu.Unlock()
+
+	p.active.Wait()
+}

common/lua/program.go

@@ -0,0 +1,56 @@
+// Package lua provides shared GopherLua programs and state management for Xray scripts.
+package lua
+
+import (
+	"bufio"
+	"context"
+	"os"
+
+	glua "github.com/yuin/gopher-lua"
+	"github.com/yuin/gopher-lua/parse"
+)
+
+// Program holds immutable bytecode that can be run by independent LStates.
+type Program struct {
+	proto *glua.FunctionProto
+}
+
+// CompileFile reads and compiles a Lua file once.
+func CompileFile(path string) (*Program, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	chunk, err := parse.Parse(bufio.NewReader(f), path)
+	if err != nil {
+		return nil, err
+	}
+	proto, err := glua.Compile(chunk, path)
+	if err != nil {
+		return nil, err
+	}
+	return &Program{proto: proto}, nil
+}
+
+// NewState creates a VM, makes modules available, and executes the file top level.
+// Module loaders run only when Lua calls require. Each state gets its own globals.
+// The caller owns the returned state.
+func (p *Program) NewState(ctx context.Context, register func(*glua.LState)) (*glua.LState, error) {
+	L := glua.NewState()
+	if register != nil {
+		register(L)
+	}
+	L.SetContext(ctx)
+	L.Push(L.NewFunctionFromProto(p.proto))
+	err := L.PCall(0, 0, nil)
+	L.RemoveContext()
+	if err == nil {
+		err = ctx.Err()
+	}
+	if err != nil {
+		L.Close()
+		return nil, err
+	}
+	return L, nil
+}