dns: preserve Lua states after recoverable query errors & refactor

This commit is contained in:
Meo597
2026-10-05 06:15:49 +08:00
parent 54610a39a5
commit 248cb666f3
4 changed files with 207 additions and 133 deletions
+95 -12
View File
@@ -2,6 +2,7 @@ package dns
import (
"context"
go_errors "errors"
"os"
"path/filepath"
"strings"
@@ -12,21 +13,25 @@ import (
featureDNS "github.com/xtls/xray-core/features/dns"
)
type geoIPScriptNameServer struct {
type scriptNameServer struct {
name string
answers map[string]net.IP
errors map[string]error
ttl uint32
calls int
}
func (s *geoIPScriptNameServer) Name() string { return s.name }
func (s *geoIPScriptNameServer) IsDisableCache() bool { return true }
func (s *scriptNameServer) Name() string { return s.name }
func (s *scriptNameServer) IsDisableCache() bool { return true }
func (s *geoIPScriptNameServer) QueryIP(ctx context.Context, domain string, _ featureDNS.IPOption) ([]net.IP, uint32, error) {
func (s *scriptNameServer) QueryIP(ctx context.Context, domain string, _ featureDNS.IPOption) ([]net.IP, uint32, error) {
if err := ctx.Err(); err != nil {
return nil, 0, err
}
s.calls++
if err := s.errors[domain]; err != nil {
return nil, 0, err
}
ip, ok := s.answers[domain]
if !ok {
return nil, 0, featureDNS.ErrEmptyResponse
@@ -34,6 +39,88 @@ func (s *geoIPScriptNameServer) QueryIP(ctx context.Context, domain string, _ fe
return []net.IP{ip}, s.ttl, nil
}
func TestDNSScriptQuery(t *testing.T) {
wantIP := net.ParseIP("127.0.0.1")
upstreamErr := go_errors.New("upstream failed")
for _, tc := range []struct {
name, body string
wantIPs []net.IP
wantTTL uint32
wantErr error
wantMessage string
wantCalls uint32
}{
{name: "IPs", body: `return server:Query(domain, ipv4, ipv6, fake)`, wantIPs: []net.IP{wantIP}, wantTTL: 60, wantCalls: 2},
{name: "empty result", body: `return nil, 0`, wantErr: featureDNS.ErrEmptyResponse, wantCalls: 2},
{name: "upstream error", body: `return server:Query("failed.example", ipv4, ipv6, fake)`, wantErr: upstreamErr, wantCalls: 2},
{name: "string error", body: `return nil, nil, "blocked"`, wantMessage: "blocked", wantCalls: 2},
{name: "invalid result", body: `return false, 0`, wantMessage: "native IP slice", wantCalls: 2},
{name: "execution error", body: `error("execution failed")`, wantMessage: "execution failed", wantCalls: 1},
} {
t.Run(tc.name, func(t *testing.T) {
script := `
local server = require("xray.dns").Servers[1]
local calls = 0
function HandleDNSQuery(domain, ipv4, ipv6, fake)
calls = calls + 1
if domain == "count.example" then
local ips, _, err = server:Query("good.example", ipv4, ipv6, fake)
return ips, calls, err
end
` + tc.body + `
end
`
path := filepath.Join(t.TempDir(), "query.lua")
if err := os.WriteFile(path, []byte(script), 0o600); err != nil {
t.Fatal(err)
}
option := featureDNS.IPOption{IPv4Enable: true}
upstream := &scriptNameServer{
name: "test",
answers: map[string]net.IP{"good.example": wantIP},
errors: map[string]error{"failed.example": upstreamErr},
ttl: 60,
}
server := &DNS{
ctx: context.Background(),
clients: []*Client{{server: upstream, ipOption: &option, timeoutMs: time.Second}},
}
engine, err := newScriptEngine(path, server)
if err != nil {
t.Fatal(err)
}
defer engine.close()
ips, ttl, err := engine.query("good.example", option)
switch {
case tc.wantErr != nil:
if err != tc.wantErr {
t.Fatalf("query error = %v, want original error %v", err, tc.wantErr)
}
case tc.wantMessage != "":
if err == nil || !strings.Contains(err.Error(), tc.wantMessage) {
t.Fatalf("query error = %v, want %q", err, tc.wantMessage)
}
case err != nil:
t.Fatal(err)
}
if ttl != tc.wantTTL || len(ips) != len(tc.wantIPs) {
t.Fatalf("query = %v, TTL %d; want %v, TTL %d", ips, ttl, tc.wantIPs, tc.wantTTL)
}
for i := range ips {
if !ips[i].Equal(tc.wantIPs[i]) {
t.Fatalf("IP %d = %v, want %v", i, ips[i], tc.wantIPs[i])
}
}
ips, calls, err := engine.query("count.example", option)
if err != nil || calls != tc.wantCalls || len(ips) != 1 || !ips[0].Equal(wantIP) {
t.Fatalf("next query = %v, calls %d, %v; want %v, calls %d", ips, calls, err, wantIP, tc.wantCalls)
}
})
}
}
func TestDNSScriptGeoIPFallback(t *testing.T) {
t.Setenv("xray.location.asset", filepath.Join("..", "..", "resources"))
script := `
@@ -59,7 +146,7 @@ end
t.Fatal(err)
}
primary := &geoIPScriptNameServer{
primary := &scriptNameServer{
name: "primary",
answers: map[string]net.IP{
"us.example": net.ParseIP("2001:4860:4860::8888"),
@@ -67,7 +154,7 @@ end
},
ttl: 30,
}
fallback := &geoIPScriptNameServer{
fallback := &scriptNameServer{
name: "fallback",
answers: map[string]net.IP{"other.example": net.ParseIP("9.9.9.9")},
ttl: 60,
@@ -138,7 +225,7 @@ func TestDNSScriptRejectsInvalidStartup(t *testing.T) {
}
}
func TestDNSScriptHookErrorAndFakeDNSOption(t *testing.T) {
func TestDNSScriptFakeDNSOption(t *testing.T) {
path := filepath.Join(t.TempDir(), "script.lua")
script := `
local server = require("xray.dns").Servers[1]
@@ -146,7 +233,6 @@ local log = require("xray.log")
log.Info("DNS script loaded")
function HandleDNSQuery(domain, ipv4, ipv6, fake)
log.Debug("DNS query: ", domain)
if domain == "bad.example" then error("script failure") end
local ips, ttl, err = server:Query(domain, ipv4, ipv6, fake)
if err then log.Error("DNS failed: ", err) end
return ips, ttl, err
@@ -160,7 +246,7 @@ end
if err != nil {
t.Fatal(err)
}
upstream := &geoIPScriptNameServer{
upstream := &scriptNameServer{
name: "FakeDNS",
answers: map[string]net.IP{"good.example": net.ParseIP("198.18.0.1")},
ttl: 30,
@@ -177,9 +263,6 @@ end
}
defer server.Close()
if _, _, err := server.LookupIP("bad.example", option); err == nil || !strings.Contains(err.Error(), "script failure") {
t.Fatalf("hook failure = %v, want script failure", err)
}
if _, _, err := server.LookupIP("good.example", option); err != featureDNS.ErrEmptyResponse {
t.Fatalf("FakeDNS without FakeEnable = %v, want ErrEmptyResponse", err)
}