add tests

app/dns/lua_test.go

@@ -2,7 +2,10 @@ package dns
 
 import (
 	"context"
+	go_errors "errors"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/xtls/xray-core/common/net"
 	featureDNS "github.com/xtls/xray-core/features/dns"
@@ -30,6 +33,90 @@ func TestDecodeLuaDNSResultNativeIP(t *testing.T) {
 	}
 }
 
+func TestDecodeLuaDNSResultNativeSliceCopiesIP(t *testing.T) {
+	L := lua.NewState()
+	defer L.Close()
+	original := net.ParseIP("8.8.8.8")
+	addresses := L.NewUserData()
+	addresses.Value = []net.IP{original}
+	result := L.NewTable()
+	result.RawSetString("ips", addresses)
+	result.RawSetString("ttl", lua.LNumber(45))
+	ips, ttl, err := decodeLuaDNSResult(result, featureDNS.IPOption{IPv4Enable: true})
+	if err != nil || ttl != 45 || len(ips) != 1 || !ips[0].Equal(original) {
+		t.Fatalf("decodeLuaDNSResult() = %v, TTL %d, %v", ips, ttl, err)
+	}
+	original[len(original)-1] = 9
+	if !ips[0].Equal(net.ParseIP("8.8.8.8")) {
+		t.Fatalf("decoded IP changed with input: %v", ips[0])
+	}
+}
+
+func TestDecodeLuaDNSResultValidation(t *testing.T) {
+	L := lua.NewState()
+	defer L.Close()
+	option := featureDNS.IPOption{IPv4Enable: true}
+
+	for _, tc := range []struct {
+		name   string
+		change func(*lua.LTable, *lua.LTable)
+		want   string
+	}{
+		{"fractional TTL", func(result, _ *lua.LTable) { result.RawSetString("ttl", lua.LNumber(1.5)) }, "invalid TTL"},
+		{"oversized TTL", func(result, _ *lua.LTable) { result.RawSetString("ttl", lua.LNumber(4294967296)) }, "invalid TTL"},
+		{"string address", func(_, addresses *lua.LTable) { addresses.RawSetInt(1, lua.LString("127.0.0.1")) }, "invalid address"},
+		{"missing addresses", func(result, _ *lua.LTable) { result.RawSetString("ips", lua.LString("127.0.0.1")) }, "must be an array"},
+		{"script error", func(result, _ *lua.LTable) { result.RawSetString("error", lua.LString("blocked by script")) }, "blocked by script"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			address := L.NewUserData()
+			address.Value = net.ParseIP("127.0.0.1")
+			addresses := L.NewTable()
+			addresses.RawSetInt(1, address)
+			result := L.NewTable()
+			result.RawSetString("ips", addresses)
+			result.RawSetString("ttl", lua.LNumber(60))
+			tc.change(result, addresses)
+			_, _, err := decodeLuaDNSResult(result, option)
+			if err == nil || !strings.Contains(err.Error(), tc.want) {
+				t.Fatalf("decodeLuaDNSResult error = %v, want %q", err, tc.want)
+			}
+		})
+	}
+
+	address := L.NewUserData()
+	address.Value = net.ParseIP("127.0.0.1")
+	addresses := L.NewTable()
+	addresses.RawSetInt(1, address)
+	result := L.NewTable()
+	result.RawSetString("ips", addresses)
+	result.RawSetString("ttl", lua.LNumber(60))
+	if _, _, err := decodeLuaDNSResult(result, featureDNS.IPOption{IPv6Enable: true}); err == nil {
+		t.Fatal("decodeLuaDNSResult accepted IPv4 with IPv6-only option")
+	}
+	result.RawSetString("ips", L.NewTable())
+	if _, _, err := decodeLuaDNSResult(result, option); !go_errors.Is(err, featureDNS.ErrEmptyResponse) {
+		t.Fatalf("empty result error = %v, want ErrEmptyResponse", err)
+	}
+}
+
+func TestCallLuaHookCancellation(t *testing.T) {
+	L := lua.NewState()
+	defer L.Close()
+	if err := L.DoString(`function handleDNSQuery(q) while true do end end`); err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	defer cancel()
+	_, _, err := (&DNS{}).CallLuaHook(L, ctx, "example.com", featureDNS.IPOption{IPv4Enable: true})
+	if err == nil {
+		t.Fatal("CallLuaHook did not stop after context cancellation")
+	}
+	if L.Context() != nil {
+		t.Fatal("CallLuaHook left the canceled context on the Lua state")
+	}
+}
+
 func TestCallLuaHookNormalizesDomain(t *testing.T) {
 	L := lua.NewState()
 	defer L.Close()