pull/720/merge
Darien Raymond 2017-11-19 21:43:20 +01:00
parent 2c26705a00
commit afacee9e05
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
1 changed files with 30 additions and 33 deletions

View File

@ -21,9 +21,7 @@ const (
) )
var ( var (
pseudoDestination = net.UDPDestination(net.LocalHostIP, net.Port(53)) multiQuestionDNS = map[net.Address]bool{
multiQuestionDns = map[net.Address]bool{
net.IPAddress([]byte{8, 8, 8, 8}): true, net.IPAddress([]byte{8, 8, 8, 8}): true,
net.IPAddress([]byte{8, 8, 4, 4}): true, net.IPAddress([]byte{8, 8, 4, 4}): true,
net.IPAddress([]byte{9, 9, 9, 9}): true, net.IPAddress([]byte{9, 9, 9, 9}): true,
@ -61,47 +59,47 @@ func NewUDPNameServer(address net.Destination, dispatcher dispatcher.Interface)
return s return s
} }
func (v *UDPNameServer) Cleanup() { func (s *UDPNameServer) Cleanup() {
expiredRequests := make([]uint16, 0, 16) expiredRequests := make([]uint16, 0, 16)
now := time.Now() now := time.Now()
v.Lock() s.Lock()
for id, r := range v.requests { for id, r := range s.requests {
if r.expire.Before(now) { if r.expire.Before(now) {
expiredRequests = append(expiredRequests, id) expiredRequests = append(expiredRequests, id)
close(r.response) close(r.response)
} }
} }
for _, id := range expiredRequests { for _, id := range expiredRequests {
delete(v.requests, id) delete(s.requests, id)
} }
v.Unlock() s.Unlock()
} }
func (v *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 { func (s *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
var id uint16 var id uint16
v.Lock() s.Lock()
if len(v.requests) > CleanupThreshold && v.nextCleanup.Before(time.Now()) { if len(s.requests) > CleanupThreshold && s.nextCleanup.Before(time.Now()) {
v.nextCleanup = time.Now().Add(CleanupInterval) s.nextCleanup = time.Now().Add(CleanupInterval)
go v.Cleanup() go s.Cleanup()
} }
for { for {
id = dice.RollUint16() id = dice.RollUint16()
if _, found := v.requests[id]; found { if _, found := s.requests[id]; found {
continue continue
} }
log.Trace(newError("add pending request id ", id).AtDebug()) log.Trace(newError("add pending request id ", id).AtDebug())
v.requests[id] = &PendingRequest{ s.requests[id] = &PendingRequest{
expire: time.Now().Add(time.Second * 8), expire: time.Now().Add(time.Second * 8),
response: response, response: response,
} }
break break
} }
v.Unlock() s.Unlock()
return id return id
} }
func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) { func (s *UDPNameServer) HandleResponse(payload *buf.Buffer) {
msg := new(dns.Msg) msg := new(dns.Msg)
err := msg.Unpack(payload.Bytes()) err := msg.Unpack(payload.Bytes())
if err == dns.ErrTruncated { if err == dns.ErrTruncated {
@ -117,14 +115,14 @@ func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) {
ttl := uint32(3600) // an hour ttl := uint32(3600) // an hour
log.Trace(newError("handling response for id ", id, " content: ", msg).AtDebug()) log.Trace(newError("handling response for id ", id, " content: ", msg).AtDebug())
v.Lock() s.Lock()
request, found := v.requests[id] request, found := s.requests[id]
if !found { if !found {
v.Unlock() s.Unlock()
return return
} }
delete(v.requests, id) delete(s.requests, id)
v.Unlock() s.Unlock()
for _, rr := range msg.Answer { for _, rr := range msg.Answer {
switch rr := rr.(type) { switch rr := rr.(type) {
@ -146,8 +144,7 @@ func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) {
close(request.response) close(request.response)
} }
func (v *UDPNameServer) BuildQueryA(domain string, id uint16) *buf.Buffer { func (s *UDPNameServer) BuildQueryA(domain string, id uint16) *buf.Buffer {
msg := new(dns.Msg) msg := new(dns.Msg)
msg.Id = id msg.Id = id
msg.RecursionDesired = true msg.RecursionDesired = true
@ -157,7 +154,7 @@ func (v *UDPNameServer) BuildQueryA(domain string, id uint16) *buf.Buffer {
Qtype: dns.TypeA, Qtype: dns.TypeA,
Qclass: dns.ClassINET, Qclass: dns.ClassINET,
}} }}
if multiQuestionDns[v.address.Address] { if multiQuestionDNS[s.address.Address] {
msg.Question = append(msg.Question, dns.Question{ msg.Question = append(msg.Question, dns.Question{
Name: dns.Fqdn(domain), Name: dns.Fqdn(domain),
Qtype: dns.TypeAAAA, Qtype: dns.TypeAAAA,
@ -174,23 +171,23 @@ func (v *UDPNameServer) BuildQueryA(domain string, id uint16) *buf.Buffer {
return buffer return buffer
} }
func (v *UDPNameServer) QueryA(domain string) <-chan *ARecord { func (s *UDPNameServer) QueryA(domain string) <-chan *ARecord {
response := make(chan *ARecord, 1) response := make(chan *ARecord, 1)
id := v.AssignUnusedID(response) id := s.AssignUnusedID(response)
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
v.udpServer.Dispatch(ctx, v.address, v.BuildQueryA(domain, id), v.HandleResponse) s.udpServer.Dispatch(ctx, s.address, s.BuildQueryA(domain, id), s.HandleResponse)
go func() { go func() {
for i := 0; i < 2; i++ { for i := 0; i < 2; i++ {
time.Sleep(time.Second) time.Sleep(time.Second)
v.Lock() s.Lock()
_, found := v.requests[id] _, found := s.requests[id]
v.Unlock() s.Unlock()
if !found { if !found {
break break
} }
v.udpServer.Dispatch(ctx, v.address, v.BuildQueryA(domain, id), v.HandleResponse) s.udpServer.Dispatch(ctx, s.address, s.BuildQueryA(domain, id), s.HandleResponse)
} }
cancel() cancel()
}() }()
@ -201,7 +198,7 @@ func (v *UDPNameServer) QueryA(domain string) <-chan *ARecord {
type LocalNameServer struct { type LocalNameServer struct {
} }
func (v *LocalNameServer) QueryA(domain string) <-chan *ARecord { func (*LocalNameServer) QueryA(domain string) <-chan *ARecord {
response := make(chan *ARecord, 1) response := make(chan *ARecord, 1)
go func() { go func() {