clean up dns package

This commit is contained in:
Darien Raymond committed 2017-11-15 00:36:14 +01:00
1 parent a430e2065a
commit 0dbfb66126
16 files changed
+225 -79

No files matched your search

+16 -14
View File
@@ -2,12 +2,14 @@ package server
import (
"context"
"fmt"
"sync"
"time"
"github.com/miekg/dns"
"v2ray.com/core/app/dispatcher"
"v2ray.com/core/app/log"
"v2ray.com/core/common"
"v2ray.com/core/common/buf"
"v2ray.com/core/common/dice"
"v2ray.com/core/common/net"
@@ -15,7 +17,6 @@ import (
)
const (
DefaultTTL = uint32(3600)
CleanupInterval = time.Second * 120
CleanupThreshold = 512
)
@@ -55,7 +56,6 @@ func NewUDPNameServer(address net.Destination, dispatcher dispatcher.Interface)
return s
}
// Private: Visible for testing.
func (v *UDPNameServer) Cleanup() {
expiredRequests := make([]uint16, 0, 16)
now := time.Now()
@@ -70,10 +70,8 @@ func (v *UDPNameServer) Cleanup() {
delete(v.requests, id)
}
v.Unlock()
expiredRequests = nil
}
// Private: Visible for testing.
func (v *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
var id uint16
v.Lock()
@@ -98,7 +96,6 @@ func (v *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
return id
}
// Private: Visible for testing.
func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) {
msg := new(dns.Msg)
err := msg.Unpack(payload.Bytes())
@@ -110,8 +107,8 @@ func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) {
IPs: make([]net.IP, 0, 16),
}
id := msg.Id
ttl := DefaultTTL
log.Trace(newError("handling response for id ", id, " content: ", msg.String()).AtDebug())
ttl := uint32(3600) // an hour
log.Trace(newError("handling response for id ", id, " content: ", msg).AtDebug())
v.Lock()
request, found := v.requests[id]
@@ -126,6 +123,7 @@ func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) {
switch rr := rr.(type) {
case *dns.A:
record.IPs = append(record.IPs, rr.A)
fmt.Println("Adding ans:", rr.A)
if rr.Hdr.Ttl < ttl {
ttl = rr.Hdr.Ttl
}
@@ -152,13 +150,18 @@ func (v *UDPNameServer) BuildQueryA(domain string, id uint16) *buf.Buffer {
Name: dns.Fqdn(domain),
Qtype: dns.TypeA,
Qclass: dns.ClassINET,
},
{
Name: dns.Fqdn(domain),
Qtype: dns.TypeAAAA,
Qclass: dns.ClassINET,
}}
buffer := buf.New()
buffer.AppendSupplier(func(b []byte) (int, error) {
common.Must(buffer.Reset(func(b []byte) (int, error) {
writtenBuffer, err := msg.PackBuffer(b)
return len(writtenBuffer), err
})
}))
return buffer
}
@@ -167,7 +170,7 @@ func (v *UDPNameServer) QueryA(domain string) <-chan *ARecord {
response := make(chan *ARecord, 1)
id := v.AssignUnusedID(response)
ctx, cancel := context.WithTimeout(context.Background(), time.Second*8)
ctx, cancel := context.WithCancel(context.Background())
v.udpServer.Dispatch(ctx, v.address, v.BuildQueryA(domain, id), v.HandleResponse)
go func() {
@@ -176,11 +179,10 @@ func (v *UDPNameServer) QueryA(domain string) <-chan *ARecord {
v.Lock()
_, found := v.requests[id]
v.Unlock()
if found {
v.udpServer.Dispatch(ctx, v.address, v.BuildQueryA(domain, id), v.HandleResponse)
} else {
if !found {
break
}
v.udpServer.Dispatch(ctx, v.address, v.BuildQueryA(domain, id), v.HandleResponse)
}
cancel()
}()
@@ -205,7 +207,7 @@ func (v *LocalNameServer) QueryA(domain string) <-chan *ARecord {
response <- &ARecord{
IPs: ips,
Expire: time.Now().Add(time.Second * time.Duration(DefaultTTL)),
Expire: time.Now().Add(time.Hour),
}
}()
-20
View File
@@ -1,20 +0,0 @@
package server
import (
"time"
"v2ray.com/core/common/net"
)
type IPResult struct {
IP []net.IP
TTL time.Duration
}
type Querier interface {
QueryDomain(domain string) <-chan *IPResult
}
type UDPQuerier struct {
server net.Destination
}
+34 -12
View File
@@ -21,22 +21,22 @@ const (
)
type DomainRecord struct {
A *ARecord
}
type Record struct {
IP []net.IP
Expire time.Time
LastAccess time.Time
}
func (r *Record) Expired() bool {
func (r *DomainRecord) Expired() bool {
return r.Expire.Before(time.Now())
}
func (r *DomainRecord) Inactive() bool {
now := time.Now()
return r.Expire.Before(now) || r.LastAccess.Add(time.Hour).Before(now)
return r.Expire.Before(now) || r.LastAccess.Add(time.Minute*5).Before(now)
}
type CacheServer struct {
sync.RWMutex
sync.Mutex
hosts map[string]net.IP
records map[string]*DomainRecord
servers []NameServer
@@ -90,15 +90,33 @@ func (*CacheServer) Start() error {
func (*CacheServer) Close() {}
func (s *CacheServer) GetCached(domain string) []net.IP {
s.RLock()
defer s.RUnlock()
s.Lock()
defer s.Unlock()
if record, found := s.records[domain]; found && record.A.Expire.After(time.Now()) {
return record.A.IPs
if record, found := s.records[domain]; found && !record.Expired() {
record.LastAccess = time.Now()
return record.IP
}
return nil
}
func (s *CacheServer) tryCleanup() {
s.Lock()
defer s.Unlock()
if len(s.records) > 256 {
domains := make([]string, 0, 256)
for d, r := range s.records {
if r.Expired() {
domains = append(domains, d)
}
}
for _, d := range domains {
delete(s.records, d)
}
}
}
func (s *CacheServer) Get(domain string) []net.IP {
if ip, found := s.hosts[domain]; found {
return []net.IP{ip}
@@ -110,6 +128,8 @@ func (s *CacheServer) Get(domain string) []net.IP {
return ips
}
s.tryCleanup()
for _, server := range s.servers {
response := server.QueryA(domain)
select {
@@ -119,7 +139,9 @@ func (s *CacheServer) Get(domain string) []net.IP {
}
s.Lock()
s.records[domain] = &DomainRecord{
A: a,
IP: a.IPs,
Expire: a.Expire,
LastAccess: time.Now(),
}
s.Unlock()
log.Trace(newError("returning ", len(a.IPs), " IPs for domain ", domain).AtDebug())
+102
View File
@@ -0,0 +1,102 @@
package server_test
import (
"context"
"testing"
"v2ray.com/core/app"
"v2ray.com/core/app/dispatcher"
_ "v2ray.com/core/app/dispatcher/impl"
. "v2ray.com/core/app/dns"
_ "v2ray.com/core/app/dns/server"
"v2ray.com/core/app/proxyman"
_ "v2ray.com/core/app/proxyman/outbound"
"v2ray.com/core/common"
"v2ray.com/core/common/net"
"v2ray.com/core/common/serial"
"v2ray.com/core/proxy/freedom"
"v2ray.com/core/testing/servers/udp"
. "v2ray.com/ext/assert"
"github.com/miekg/dns"
)
type staticHandler struct {
}
func (*staticHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
ans := new(dns.Msg)
ans.Id = r.Id
for _, q := range r.Question {
if q.Name == "google.com." && q.Qtype == dns.TypeA {
rr, _ := dns.NewRR("google.com. IN A 8.8.8.8")
ans.Answer = append(ans.Answer, rr)
} else if q.Name == "facebook.com." && q.Qtype == dns.TypeA {
rr, _ := dns.NewRR("facebook.com. IN A 9.9.9.9")
ans.Answer = append(ans.Answer, rr)
}
}
w.WriteMsg(ans)
}
func TestUDPServer(t *testing.T) {
assert := With(t)
port := udp.PickPort()
dnsServer := dns.Server{
Addr: "127.0.0.1:" + port.String(),
Net: "udp",
Handler: &staticHandler{},
UDPSize: 1200,
}
go dnsServer.ListenAndServe()
config := &Config{
NameServers: []*net.Endpoint{
{
Network: net.Network_UDP,
Address: &net.IPOrDomain{
Address: &net.IPOrDomain_Ip{
Ip: []byte{127, 0, 0, 1},
},
},
Port: uint32(port),
},
},
}
ctx := context.Background()
space := app.NewSpace()
ctx = app.ContextWithSpace(ctx, space)
common.Must(app.AddApplicationToSpace(ctx, config))
common.Must(app.AddApplicationToSpace(ctx, &dispatcher.Config{}))
common.Must(app.AddApplicationToSpace(ctx, &proxyman.OutboundConfig{}))
om := proxyman.OutboundHandlerManagerFromSpace(space)
om.AddHandler(ctx, &proxyman.OutboundHandlerConfig{
ProxySettings: serial.ToTypedMessage(&freedom.Config{}),
})
common.Must(space.Initialize())
common.Must(space.Start())
server := FromSpace(space)
assert(server, IsNotNil)
ips := server.Get("google.com")
assert(len(ips), Equals, 1)
assert([]byte(ips[0]), Equals, []byte{8, 8, 8, 8})
ips = server.Get("facebook.com")
assert(len(ips), Equals, 1)
assert([]byte(ips[0]), Equals, []byte{9, 9, 9, 9})
dnsServer.Shutdown()
ips = server.Get("google.com")
assert(len(ips), Equals, 1)
assert([]byte(ips[0]), Equals, []byte{8, 8, 8, 8})
}