188 lines
3.3 KiB
Go
188 lines
3.3 KiB
Go
package dnsforward
|
|
|
|
import (
|
|
"net"
|
|
"testing"
|
|
|
|
"github.com/AdguardTeam/dnsproxy/proxy"
|
|
"github.com/miekg/dns"
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
// fakeIpsetMgr is a fake aghnet.IpsetManager for tests.
|
|
type fakeIpsetMgr struct {
|
|
ip4s []net.IP
|
|
ip6s []net.IP
|
|
}
|
|
|
|
// Add implements the aghnet.IpsetManager interface for *fakeIpsetMgr.
|
|
func (m *fakeIpsetMgr) Add(host string, ip4s, ip6s []net.IP) (n int, err error) {
|
|
m.ip4s = append(m.ip4s, ip4s...)
|
|
m.ip6s = append(m.ip6s, ip6s...)
|
|
|
|
return len(ip4s) + len(ip6s), nil
|
|
}
|
|
|
|
// Close implements the aghnet.IpsetManager interface for *fakeIpsetMgr.
|
|
func (*fakeIpsetMgr) Close() (err error) {
|
|
return nil
|
|
}
|
|
|
|
func TestIpsetCtx_process(t *testing.T) {
|
|
ip4 := net.IP{1, 2, 3, 4}
|
|
ip6 := net.IP{
|
|
0x12, 0x34, 0x00, 0x00,
|
|
0x00, 0x00, 0x00, 0x00,
|
|
0x00, 0x00, 0x00, 0x00,
|
|
0x00, 0x00, 0x56, 0x78,
|
|
}
|
|
|
|
req4 := createTestMessageWithType("example.com", dns.TypeA)
|
|
req6 := createTestMessageWithType("example.com", dns.TypeAAAA)
|
|
|
|
resp4 := &dns.Msg{
|
|
Answer: []dns.RR{&dns.A{
|
|
A: ip4,
|
|
}},
|
|
}
|
|
resp6 := &dns.Msg{
|
|
Answer: []dns.RR{&dns.AAAA{
|
|
AAAA: ip6,
|
|
}},
|
|
}
|
|
|
|
t.Run("nil", func(t *testing.T) {
|
|
dctx := &dnsContext{
|
|
proxyCtx: &proxy.DNSContext{},
|
|
|
|
responseFromUpstream: true,
|
|
}
|
|
|
|
ictx := &ipsetCtx{}
|
|
rc := ictx.process(dctx)
|
|
assert.Equal(t, resultCodeSuccess, rc)
|
|
|
|
err := ictx.close()
|
|
assert.NoError(t, err)
|
|
})
|
|
|
|
t.Run("ipv4", func(t *testing.T) {
|
|
dctx := &dnsContext{
|
|
proxyCtx: &proxy.DNSContext{
|
|
Req: req4,
|
|
Res: resp4,
|
|
},
|
|
|
|
responseFromUpstream: true,
|
|
}
|
|
|
|
m := &fakeIpsetMgr{}
|
|
ictx := &ipsetCtx{
|
|
ipsetMgr: m,
|
|
}
|
|
|
|
rc := ictx.process(dctx)
|
|
assert.Equal(t, resultCodeSuccess, rc)
|
|
assert.Equal(t, []net.IP{ip4}, m.ip4s)
|
|
assert.Empty(t, m.ip6s)
|
|
|
|
err := ictx.close()
|
|
assert.NoError(t, err)
|
|
})
|
|
|
|
t.Run("ipv6", func(t *testing.T) {
|
|
dctx := &dnsContext{
|
|
proxyCtx: &proxy.DNSContext{
|
|
Req: req6,
|
|
Res: resp6,
|
|
},
|
|
|
|
responseFromUpstream: true,
|
|
}
|
|
|
|
m := &fakeIpsetMgr{}
|
|
ictx := &ipsetCtx{
|
|
ipsetMgr: m,
|
|
}
|
|
|
|
rc := ictx.process(dctx)
|
|
assert.Equal(t, resultCodeSuccess, rc)
|
|
assert.Empty(t, m.ip4s)
|
|
assert.Equal(t, []net.IP{ip6}, m.ip6s)
|
|
|
|
err := ictx.close()
|
|
assert.NoError(t, err)
|
|
})
|
|
}
|
|
|
|
func TestIpsetCtx_SkipIpsetProcessing(t *testing.T) {
|
|
req4 := createTestMessage("example.com")
|
|
resp4 := &dns.Msg{
|
|
Answer: []dns.RR{&dns.A{
|
|
A: net.IP{1, 2, 3, 4},
|
|
}},
|
|
}
|
|
|
|
m := &fakeIpsetMgr{}
|
|
ictx := &ipsetCtx{
|
|
ipsetMgr: m,
|
|
}
|
|
|
|
testCases := []struct {
|
|
dctx *dnsContext
|
|
name string
|
|
want bool
|
|
}{{
|
|
name: "basic",
|
|
want: false,
|
|
dctx: &dnsContext{
|
|
proxyCtx: &proxy.DNSContext{
|
|
Req: req4,
|
|
Res: resp4,
|
|
},
|
|
|
|
responseFromUpstream: true,
|
|
},
|
|
}, {
|
|
name: "rewrite",
|
|
want: true,
|
|
dctx: &dnsContext{
|
|
proxyCtx: &proxy.DNSContext{
|
|
Req: req4,
|
|
Res: resp4,
|
|
},
|
|
|
|
responseFromUpstream: false,
|
|
},
|
|
}, {
|
|
name: "empty_req",
|
|
want: true,
|
|
dctx: &dnsContext{
|
|
proxyCtx: &proxy.DNSContext{
|
|
Req: nil,
|
|
Res: resp4,
|
|
},
|
|
|
|
responseFromUpstream: true,
|
|
},
|
|
}, {
|
|
name: "empty_res",
|
|
want: true,
|
|
dctx: &dnsContext{
|
|
proxyCtx: &proxy.DNSContext{
|
|
Req: req4,
|
|
Res: nil,
|
|
},
|
|
|
|
responseFromUpstream: true,
|
|
},
|
|
}}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
got := ictx.skipIpsetProcessing(tc.dctx)
|
|
assert.Equal(t, tc.want, got)
|
|
})
|
|
}
|
|
}
|